class_docstr: Discretize a continuous relaxation of a one-hot encoded categorical.
  
  This returns a one-hot encoded categorical and use a straight-through
  gradient estimator via an identity function.
  
  [Daulton2022bopr]_ proposes using this in acquisition optimization.
class_docstr_tok: ["Discretize","a","continuous","relaxation","of","a","one-hot","encoded","categorical",".","This","returns","a","one-hot","encoded","categorical","and","use","a","straight-through","gradient","estimator","via","an","identity","function",".","[","Daulton2022bopr","]","_","proposes","using","this","in","acquisition","optimization","."]
docstr: Discretize the input tensor.
  
  This applies a argmax along the last dimensions of the input tensor
  and one-hot encodes the result.
  
  Args:
      X: The tensor to be rounded.
  
  Returns:
      A tensor where each element is rounded to the nearest integer.
docstr_tok: ["Discretize","the","input","tensor",".","This","applies","a","argmax","along","the","last","dimensions","of","the","input","tensor","and","one-hot","encodes","the","result",".","Args",":","X",":","The","tensor","to","be","rounded",".","Returns",":","A","tensor","where","each","element","is","rounded","to","the","nearest","integer","."]
code: def forward(ctx, X: Tensor) -> Tensor:
      r"""Discretize the input tensor.
  
      This applies a argmax along the last dimensions of the input tensor
      and one-hot encodes the result.
  
      Args:
          X: The tensor to be rounded.
  
      Returns:
          A tensor where each element is rounded to the nearest integer.
      """
      return one_hot(X.argmax(dim=-1), num_classes=X.shape[-1]).to(X)
  
file_path: botorch/botorch/utils/rounding.py
imports: from __future__ import annotations
  import torch
  from torch import Tensor
  from torch.autograd import Function
  from torch.nn.functional import one_hot
label_desc: Use image node_id 1 for calling the OneHotArgmaxSTE obj's underlying member method code with example usage: obj.forward(ctx, X) and returns: one_hot