def _one_hot_tensor_3d(x, index, total_length):
  """One-hot encodes a 2d Tensor in a 3d Tensor.

  This could potentially be implemented in a simpler way using tf.pad but
  this method is compatible with XLA's restriction on static shapes defined
  by constants.

  Args:
    x: A Tensor of shape [m, n].
    index: The "on" index of the inner dimension of the output Tensor.
    total_length: Total length of the output Tensor.

  Returns:
    A Tensor of shape [m, n, total_length].
  """
  m = x.get_shape()[0]
  n = x.get_shape()[1]
  x = tf.expand_dims(x, 2)
  index_tiled = tf.tile(tf.expand_dims(index, 0), [m * n])
  one_hot = tf.one_hot(index_tiled, total_length, on_value=1)
  one_hot = tf.cast(one_hot, dtype=x.dtype)
  one_hot = tf.reshape(one_hot, [m, n, total_length])
  return one_hot * x
One-hot encodes a 2d Tensor in a 3d Tensor.

  This could potentially be implemented in a simpler way using tf.pad but
  this method is compatible with XLA's restriction on static shapes defined
  by constants.

  Args:
    x: A Tensor of shape [m, n].
    index: The "on" index of the inner dimension of the output Tensor.
    total_length: Total length of the output Tensor.

  Returns:
    A Tensor of shape [m, n, total_length].
  
language/xsp/model/beam_search.py