jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def linspace(
    start, stop, num=50, endpoint=True, retstep=False, dtype=None, axis=0
):
    if dtype is None:
        dtypes_to_resolve = [
            getattr(start, "dtype", type(start)),
            getattr(stop, "dtype", type(stop)),
            float,
        ]
        dtype = dtypes.result_type(*dtypes_to_resolve)
    else:
        dtype = standardize_dtype(dtype)
    start = convert_to_tensor(start, dtype=dtype)
    stop = convert_to_tensor(stop, dtype=dtype)
    if num < 0:
        raise ValueError(
            f"`num` must be a non-negative integer. Received: num={num}"
        )
    step = tf.convert_to_tensor(np.nan)
    if endpoint:
        result = tf.linspace(start, stop, num, axis=axis)
        if num > 1:
            step = (stop - start) / (num - 1)
    else:
        # tf.linspace doesn't support endpoint=False, so we manually handle it
        if num > 0:
            step = (stop - start) / num
        if num > 1:
            new_stop = tf.cast(stop, step.dtype) - step
            start = tf.cast(start, new_stop.dtype)
            result = tf.linspace(start, new_stop, num, axis=axis)
        else:
            result = tf.linspace(start, stop, num, axis=axis)
    if dtype is not None:
        if "int" in dtype:
            result = tf.floor(result)
        result = tf.cast(result, dtype)
    if retstep:
        return (result, step)
    else:
        return result
//  [python]
def logspace(start, stop, num=50, endpoint=True, base=10, dtype=None, axis=0):
    result = linspace(
        start=start,
        stop=stop,
        num=num,
        endpoint=endpoint,
        dtype=dtype,
        axis=axis,
    )
    return tf.pow(t.cast(base, result.dtype), result)
//  [python]
def logspace(start, stop, num=50, endpoint=True, base=10, dtype=None, axis=0):
    result = linspace(
        start=start,
        stop=stop,
        num=num,
        endpoint=endpoint,
        dtype=dtype,
        axis=axis,
    )
    return tf.pow(tf.cast(base, result.dtype), result)
//  [python]
def nonzero(x):
    x = convert_to_tensor(x)
    result = tf.unstack(tf.where(tf.cast(x, "bool")), x.shape.rank, axis=1)
    return tree.map_structure(lambda indices: t.cast(indices, "int32"), result)
//  [python]
def nonzero(x):
    x = convert_to_tensor(x)
    result = tf.unstack(tf.where(tf.cast(x, "bool")), x.shape.rank, axis=1)
    return tree.map_structure(lambda indices: tf.cast(indices, "int32"), result)
//  [python]
def mod(x1, x2):
    x1 = convert_to_tensor(x1)
    x2 = convert_to_tensor(x2)
    dtype = dtypes.result_type(x1.dtype, x2.dtype)
    if dtype == "bool":
        dtype = "int32"
    x1 = tf.cast(x1, dtype)
    x2 = tf.cast(x2, dtype)
    return tf.math.mod(x1, x2)
//  [python]
def not_equal(x1, x2):
    x1 = convert_to_tensor(x1)
    x2 = convert_to_tensor(x2)
    dtype = dtypes.result_type(x1.dtype, x2.dtype)
    x1 = tf.cast(x1, dtype)
    x2 = tf.cast(x2, dtype)
    return tf.not_equal(x1, x2)
//  [python]
def outer(x1, x2):
    x1 = convert_to_tensor(x1)
    x2 = convert_to_tensor(x2)
    dtype = dtypes.result_type(x1.dtype, x2.dtype)
    x1 = t.cast(x1, dtype)
    x2 = t.cast(x2, dtype)
    return tf.reshape(x1, [-1, 1]) * tf.reshape(x2, [-1])
//  [python]
def outer(x1, x2):
    x1 = convert_to_tensor(x1)
    x2 = convert_to_tensor(x2)
    dtype = dtypes.result_type(x1.dtype, x2.dtype)
    x1 = tf.cast(x1, dtype)
    x2 = tf.cast(x2, dtype)
    return tf.reshape(x1, [-1, 1]) * tf.reshape(x2, [-1])
//  [python]
def min(x, axis=None, keepdims=False, initial=None):
    x = convert_to_tensor(x)

    # The TensorFlow numpy API implementation doesn't support `initial` so we
    # handle it manually here.
    if initial is not None:
        if standardize_dtype(x.dtype) == "bool":
            x = tf.reduce_all(x, axis=axis, keepdims=keepdims)
            x = tf.math.minimum(t.cast(x, "int32"), t.cast(initial, "int32"))
            return t.cast(x, "bool")
        else:
            x = tf.reduce_min(x, axis=axis, keepdims=keepdims)
        return tf.math.minimum(x, initial)

    # TensorFlow returns inf by default for an empty list, but for consistency
    # with other backends and the numpy API we want to throw in this case.
    if tf.executing_eagerly():
        size_x = size(x)
        tf.assert_greater(
            size_x,
            tf.constant(0, dtype=size_x.dtype),
            message="Cannot compute the min of an empty tensor.",
        )

    if standardize_dtype(x.dtype) == "bool":
        return tf.reduce_all(x, axis=axis, keepdims=keepdims)
    else:
        return tf.reduce_min(x, axis=axis, keepdims=keepdims)
//  [python]
def min(x, axis=None, keepdims=False, initial=None):
    x = convert_to_tensor(x)

    # The TensorFlow numpy API implementation doesn't support `initial` so we
    # handle it manually here.
    if initial is not None:
        if standardize_dtype(x.dtype) == "bool":
            x = tf.reduce_all(x, axis=axis, keepdims=keepdims)
            x = tf.math.minimum(tf.cast(x, "int32"), tf.cast(initial, "int32"))
            return tf.cast(x, "bool")
        else:
            x = tf.reduce_min(x, axis=axis, keepdims=keepdims)
        return tf.math.minimum(x, initial)

    # TensorFlow returns inf by default for an empty list, but for consistency
    # with other backends and the numpy API we want to throw in this case.
    if tf.executing_eagerly():
        size_x = size(x)
        tf.assert_greater(
            size_x,
            tf.constant(0, dtype=size_x.dtype),
            message="Cannot compute the min of an empty tensor.",
        )

    if standardize_dtype(x.dtype) == "bool":
        return tf.reduce_all(x, axis=axis, keepdims=keepdims)
    else:
        return tf.reduce_min(x, axis=axis, keepdims=keepdims)
//  [python]
def prod(x, axis=None, keepdims=False, dtype=None):
    x = convert_to_tensor(x)
    if dtype is None:
        dtype = dtypes.result_type(x.dtype)
        if dtype == "bool":
            dtype = "int32"
        elif dtype in ("int8", "int16"):
            dtype = "int32"
        elif dtype in ("uint8", "uint16"):
            dtype = "uint32"
        x = t.cast(x, dtype)
    return tf.reduce_prod(x, axis=axis, keepdims=keepdims)
//  [python]
def prod(x, axis=None, keepdims=False, dtype=None):
    x = convert_to_tensor(x)
    if dtype is None:
        dtype = dtypes.result_type(x.dtype)
        if dtype == "bool":
            dtype = "int32"
        elif dtype in ("int8", "int16"):
            dtype = "int32"
        elif dtype in ("uint8", "uint16"):
            dtype = "uint32"
        x = tf.cast(x, dtype)
    return tf.reduce_prod(x, axis=axis, keepdims=keepdims)
//  [python]
def repeat(x, repeats, axis=None):
    x = convert_to_tensor(x)
    # TODO: tf.repeat doesn't support uint16
    if standardize_dtype(x.dtype) == "uint16":
        x = t.cast(x, "uint32")
        return t.cast(tf.repeat(x, repeats, axis=axis), "uint16")
    return tf.repeat(x, repeats, axis=axis)
//  [python]
def repeat(x, repeats, axis=None):
    x = convert_to_tensor(x)
    # TODO: tf.repeat doesn't support uint16
    if standardize_dtype(x.dtype) == "uint16":
        x = tf.cast(x, "uint32")
        return tf.cast(tf.repeat(x, repeats, axis=axis), "uint16")
    return tf.repeat(x, repeats, axis=axis)
//  [python]
def quantile(x, q, axis=None, method="linear", keepdims=False):
    x = convert_to_tensor(x)
    q = convert_to_tensor(q)
    axis = to_tuple_or_list(axis)
    compute_dtype = dtypes.result_type(x.dtype, float)
    x = t.cast(x, compute_dtype)
    return _quantile(x, q, axis=axis, method=method, keepdims=keepdims)
//  [python]
def quantile(x, q, axis=None, method="linear", keepdims=False):
    x = convert_to_tensor(x)
    q = convert_to_tensor(q)
    axis = to_tuple_or_list(axis)
    compute_dtype = dtypes.result_type(x.dtype, float)
    x = tf.cast(x, compute_dtype)
    return _quantile(x, q, axis=axis, method=method, keepdims=keepdims)
//  [python]
def sign(x):
    x = convert_to_tensor(x)
    ori_dtype = standardize_dtype(x.dtype)
    # TODO: tf.sign doesn't support uint8, uint16, uint32
    if ori_dtype in ("uint8", "uint16", "uint32"):
        x = t.cast(x, "int32")
        return t.cast(tf.sign(x), ori_dtype)
    return tf.sign(x)
//  [python]
def sign(x):
    x = convert_to_tensor(x)
    ori_dtype = standardize_dtype(x.dtype)
    # TODO: tf.sign doesn't support uint8, uint16, uint32
    if ori_dtype in ("uint8", "uint16", "uint32"):
        x = tf.cast(x, "int32")
        return tf.cast(tf.sign(x), ori_dtype)
    return tf.sign(x)
//  [python]
def sin(x):
    x = convert_to_tensor(x)
    if standardize_dtype(x.dtype) == "int64":
        dtype = config.floatx()
    else:
        dtype = dtypes.result_type(x.dtype, float)
    x = t.cast(x, dtype)
    return tf.math.sin(x)
//  [python]
def sin(x):
    x = convert_to_tensor(x)
    if standardize_dtype(x.dtype) == "int64":
        dtype = config.floatx()
    else:
        dtype = dtypes.result_type(x.dtype, float)
    x = tf.cast(x, dtype)
    return tf.math.sin(x)
//  [python]
def sinh(x):
    x = convert_to_tensor(x)
    if standardize_dtype(x.dtype) == "int64":
        dtype = config.floatx()
    else:
        dtype = dtypes.result_type(x.dtype, float)
    x = t.cast(x, dtype)
    return tf.math.sinh(x)
//  [python]
def sinh(x):
    x = convert_to_tensor(x)
    if standardize_dtype(x.dtype) == "int64":
        dtype = config.floatx()
    else:
        dtype = dtypes.result_type(x.dtype, float)
    x = tf.cast(x, dtype)
    return tf.math.sinh(x)
//  [python]
def sort(x, axis=-1):
    x = convert_to_tensor(x)
    ori_dtype = standardize_dtype(x.dtype)
    # TODO: tf.sort doesn't support bool
    if ori_dtype == "bool":
        x = t.cast(x, "int8")
        return t.cast(tf.sort(x, axis=axis), ori_dtype)
    return tf.sort(x, axis=axis)
//  [python]
def sort(x, axis=-1):
    x = convert_to_tensor(x)
    ori_dtype = standardize_dtype(x.dtype)
    # TODO: tf.sort doesn't support bool
    if ori_dtype == "bool":
        x = tf.cast(x, "int8")
        return tf.cast(tf.sort(x, axis=axis), ori_dtype)
    return tf.sort(x, axis=axis)
//  [python]
def std(x, axis=None, keepdims=False):
    x = convert_to_tensor(x)
    ori_dtype = standardize_dtype(x.dtype)
    if "int" in ori_dtype or ori_dtype == "bool":
        x = t.cast(x, config.floatx())
    return tf.math.reduce_std(x, axis=axis, keepdims=keepdims)
//  [python]
def std(x, axis=None, keepdims=False):
    x = convert_to_tensor(x)
    ori_dtype = standardize_dtype(x.dtype)
    if "int" in ori_dtype or ori_dtype == "bool":
        x = tf.cast(x, config.floatx())
    return tf.math.reduce_std(x, axis=axis, keepdims=keepdims)
//  [python]
def tan(x):
    x = convert_to_tensor(x)
    if standardize_dtype(x.dtype) == "int64":
        dtype = config.floatx()
    else:
        dtype = dtypes.result_type(x.dtype, float)
    x = t.cast(x, dtype)
    return tf.math.tan(x)
//  [python]
def tan(x):
    x = convert_to_tensor(x)
    if standardize_dtype(x.dtype) == "int64":
        dtype = config.floatx()
    else:
        dtype = dtypes.result_type(x.dtype, float)
    x = tf.cast(x, dtype)
    return tf.math.tan(x)
//  [python]
def tanh(x):
    x = convert_to_tensor(x)
    if standardize_dtype(x.dtype) == "int64":
        dtype = config.floatx()
    else:
        dtype = dtypes.result_type(x.dtype, float)
    x = t.cast(x, dtype)
    return tf.math.tanh(x)
//  [python]
def tanh(x):
    x = convert_to_tensor(x)
    if standardize_dtype(x.dtype) == "int64":
        dtype = config.floatx()
    else:
        dtype = dtypes.result_type(x.dtype, float)
    x = tf.cast(x, dtype)
    return tf.math.tanh(x)
//  [python]
def tensordot(x1, x2, axes=2):
    x1 = convert_to_tensor(x1)
    x2 = convert_to_tensor(x2)
    result_dtype = dtypes.result_type(x1.dtype, x2.dtype)
    # TODO: tf.tensordot only supports float types
    compute_dtype = dtypes.result_type(result_dtype, float)
    x1 = t.cast(x1, compute_dtype)
    x2 = t.cast(x2, compute_dtype)
    return t.cast(tf.tensordot(x1, x2, axes=axes), dtype=result_dtype)
//  [python]
def tensordot(x1, x2, axes=2):
    x1 = convert_to_tensor(x1)
    x2 = convert_to_tensor(x2)
    result_dtype = dtypes.result_type(x1.dtype, x2.dtype)
    # TODO: tf.tensordot only supports float types
    compute_dtype = dtypes.result_type(result_dtype, float)
    x1 = tf.cast(x1, compute_dtype)
    x2 = tf.cast(x2, compute_dtype)
    return tf.cast(tf.tensordot(x1, x2, axes=axes), dtype=result_dtype)
//  [python]
def vdot(x1, x2):
    x1 = convert_to_tensor(x1)
    x2 = convert_to_tensor(x2)
    result_dtype = dtypes.result_type(x1.dtype, x2.dtype)
    compute_dtype = dtypes.result_type(result_dtype, float)
    x1 = tf.cast(x1, compute_dtype)
    x2 = tf.cast(x2, compute_dtype)
    x1 = tf.reshape(x1, [-1])
    x2 = tf.reshape(x2, [-1])
    return tf.cast(dot(x1, x2), result_dtype)
//  [python]
def square(x):
    x = convert_to_tensor(x)
    if standardize_dtype(x.dtype) == "bool":
        x = t.cast(x, "int32")
    return tf.square(x)
//  [python]
def square(x):
    x = convert_to_tensor(x)
    if standardize_dtype(x.dtype) == "bool":
        x = tf.cast(x, "int32")
    return tf.square(x)
//  [python]
def tri(N, M=None, k=0, dtype=None):
    M = M if M is not None else N
    dtype = standardize_dtype(dtype or config.floatx())
    if k < 0:
        lower = -k - 1
        if lower > N:
            r = tf.zeros([N, M], dtype=dtype)
        else:
            o = tf.ones([N, M], dtype="bool")
            r = t.cast(
                tf.logical_not(tf.linalg.band_part(o, lower, -1)), dtype=dtype
            )
    else:
        o = tf.ones([N, M], dtype=dtype)
        if k > M:
            r = o
        else:
            r = tf.linalg.band_part(o, -1, k)
    return r
//  [python]
def tri(N, M=None, k=0, dtype=None):
    M = M if M is not None else N
    dtype = standardize_dtype(dtype or config.floatx())
    if k < 0:
        lower = -k - 1
        if lower > N:
            r = tf.zeros([N, M], dtype=dtype)
        else:
            o = tf.ones([N, M], dtype="bool")
            r = tf.cast(
                tf.logical_not(tf.linalg.band_part(o, lower, -1)), dtype=dtype
            )
    else:
        o = tf.ones([N, M], dtype=dtype)
        if k > M:
            r = o
        else:
            r = tf.linalg.band_part(o, -1, k)
    return r
//  [python]
def trace(x, offset=0, axis1=0, axis2=1):
    x = convert_to_tensor(x)
    dtype = standardize_dtype(x.dtype)
    if dtype not in ("int64", "uint32", "uint64"):
        dtype = dtypes.result_type(dtype, "int32")
    x_shape = tf.shape(x)
    x = moveaxis(x, (axis1, axis2), (-2, -1))
    # Mask out the diagonal and reduce.
    x = tf.where(
        eye(x_shape[axis1], x_shape[axis2], k=offset, dtype="bool"),
        x,
        tf.zeros_like(x),
    )
    # The output dtype is set to "int32" if the input dtype is "bool"
    if standardize_dtype(x.dtype) == "bool":
        x = t.cast(x, "int32")
    return t.cast(tf.reduce_sum(x, axis=(-2, -1)), dtype)
//  [python]
def trace(x, offset=0, axis1=0, axis2=1):
    x = convert_to_tensor(x)
    dtype = standardize_dtype(x.dtype)
    if dtype not in ("int64", "uint32", "uint64"):
        dtype = dtypes.result_type(dtype, "int32")
    x_shape = tf.shape(x)
    x = moveaxis(x, (axis1, axis2), (-2, -1))
    # Mask out the diagonal and reduce.
    x = tf.where(
        eye(x_shape[axis1], x_shape[axis2], k=offset, dtype="bool"),
        x,
        tf.zeros_like(x),
    )
    # The output dtype is set to "int32" if the input dtype is "bool"
    if standardize_dtype(x.dtype) == "bool":
        x = tf.cast(x, "int32")
    return tf.cast(tf.reduce_sum(x, axis=(-2, -1)), dtype)
//  [python]
def where(condition, x1, x2):
    condition = t.cast(condition, "bool")
    if x1 is not None and x2 is not None:
        if not isinstance(x1, (int, float)):
            x1 = convert_to_tensor(x1)
        if not isinstance(x2, (int, float)):
            x2 = convert_to_tensor(x2)
        dtype = dtypes.result_type(
            getattr(x1, "dtype", type(x1)),
            getattr(x2, "dtype", type(x2)),
        )
        x1 = convert_to_tensor(x1, dtype)
        x2 = convert_to_tensor(x2, dtype)
        return tf.where(condition, x1, x2)
    if x1 is None and x2 is None:
        return nonzero(condition)
    raise ValueError(
        "`x1` and `x2` either both should be `None`"
        " or both should have non-None value."
    )
//  [python]
def where(condition, x1, x2):
    condition = tf.cast(condition, "bool")
    if x1 is not None and x2 is not None:
        if not isinstance(x1, (int, float)):
            x1 = convert_to_tensor(x1)
        if not isinstance(x2, (int, float)):
            x2 = convert_to_tensor(x2)
        dtype = dtypes.result_type(
            getattr(x1, "dtype", type(x1)),
            getattr(x2, "dtype", type(x2)),
        )
        x1 = convert_to_tensor(x1, dtype)
        x2 = convert_to_tensor(x2, dtype)
        return tf.where(condition, x1, x2)
    if x1 is None and x2 is None:
        return nonzero(condition)
    raise ValueError(
        "`x1` and `x2` either both should be `None`"
        " or both should have non-None value."
    )
//  [python]
def power(x1, x2):
    if not isinstance(x1, (int, float)):
        x1 = convert_to_tensor(x1)
    if not isinstance(x2, (int, float)):
        x2 = convert_to_tensor(x2)
    dtype = dtypes.result_type(
        getattr(x1, "dtype", type(x1)),
        getattr(x2, "dtype", type(x2)),
    )
    # TODO: tf.pow doesn't support uint* types
    if "uint" in dtype:
        x1 = convert_to_tensor(x1, "int32")
        x2 = convert_to_tensor(x2, "int32")
        return tf.cast(tf.pow(x1, x2), dtype)
    x1 = convert_to_tensor(x1, dtype)
    x2 = convert_to_tensor(x2, dtype)
    return tf.pow(x1, x2)
//  [python]
def logical_xor(x1, x2):
    x1 = t.cast(x1, "bool")
    x2 = t.cast(x2, "bool")
    return tf.math.logical_xor(x1, x2)
//  [python]
def logical_xor(x1, x2):
    x1 = tf.cast(x1, "bool")
    x2 = tf.cast(x2, "bool")
    return tf.math.logical_xor(x1, x2)
//  [python]
def sqrt(x):
    x = convert_to_tensor(x)
    dtype = (
        config.floatx()
        if standardize_dtype(x.dtype) == "int64"
        else dtypes.result_type(x.dtype, float)
    )
    x = tf.cast(x, dtype)
    return tf.math.sqrt(x)
//  [python]
def var(x, axis=None, keepdims=False):
    x = convert_to_tensor(x)
    compute_dtype = dtypes.result_type(x.dtype, "float32")
    result_dtype = dtypes.result_type(x.dtype, float)
    x = t.cast(x, compute_dtype)
    return t.cast(
        tf.math.reduce_variance(x, axis=axis, keepdims=keepdims),
        result_dtype,
    )
//  [python]
def var(x, axis=None, keepdims=False):
    x = convert_to_tensor(x)
    compute_dtype = dtypes.result_type(x.dtype, "float32")
    result_dtype = dtypes.result_type(x.dtype, float)
    x = tf.cast(x, compute_dtype)
    return tf.cast(
        tf.math.reduce_variance(x, axis=axis, keepdims=keepdims),
        result_dtype,
    )
//  [python]
def _preprocess_image(self, image):
    image = video_input.process_image(
        image=image,
        is_training=False,
        num_frames=self._num_frames,
        stride=self._stride,
        num_test_clips=1,
        min_resize=self._min_resize,
        crop_size=self._crop_size,
        num_crops=1)
    image = tf.cast(image, tf.float32)  # Use config.
    features = {'image': image}
    return features
//  [python]
def _preprocess_audio(self, audio):
    features = {}
    audio = t.cast(audio, dtype=t.float32)  # Use config.
    audio = video_input.preprocess_ops_3d.sample_sequence(
        audio, 20, random=False, stride=1)
    audio = tf.ensure_shape(
        audio, self._params.task.validation_data.audio_feature_shape)
    features['audio'] = audio
    return features
//  [python]
def _preprocess_audio(self, audio):
    features = {}
    audio = tf.cast(audio, dtype=tf.float32)  # Use config.
    audio = video_input.preprocess_ops_3d.sample_sequence(
        audio, 20, random=False, stride=1)
    audio = tf.ensure_shape(
        audio, self._params.task.validation_data.audio_feature_shape)
    features['audio'] = audio
    return features
//  [python]
def call(self, inputs, step_type=None, network_state=()):
    del step_type
    hidden_state = t.cast(tf.nest.flatten(inputs), tf.float32)[0]
    batch_squash = network_utils.BatchSquash(self._outer_rank)
    hidden_state = batch_squash.flatten(hidden_state)
    for layer in self._dummy_layers:
      hidden_state = layer(hidden_state)
    value_pred = tf.squeeze(batch_squash.unflatten(hidden_state), axis=-1)
    return value_pred, network_state
//  [python]
def call(self, inputs, step_type=None, network_state=()):
    del step_type
    hidden_state = tf.cast(tf.nest.flatten(inputs), tf.float32)[0]
    batch_squash = network_utils.BatchSquash(self._outer_rank)
    hidden_state = batch_squash.flatten(hidden_state)
    for layer in self._dummy_layers:
      hidden_state = layer(hidden_state)
    value_pred = tf.squeeze(batch_squash.unflatten(hidden_state), axis=-1)
    return value_pred, network_state
//  [python]
def _build_regressor(x, z_dim=64):
  """Make the GILBO regressor, which is based off of the GAN discriminator."""
  net = t.cast(x, tf.float32)
  net = layers.conv2d(net, 64, 4, 2, activation=tf.nn.leaky_relu)
  net = layers.conv2d(net, 128, 4, 2, activation=tf.nn.leaky_relu)
  net = layers.flatten(net)
  net = layers.dense(net, 1024, activation=tf.nn.leaky_relu)
  net = layers.dense(net, 2 * z_dim)
  # a and b correspond to the alpha beta parameters of the Beta distribution.
  a, b = net[..., :z_dim], net[..., z_dim:2 * z_dim]
  a = 1 + tf.nn.softplus(a - 5)
  b = 1 + tf.nn.softplus(b - 5)
  dist = ds.Independent(ds.Beta(a, b), 1)
  bijector = ds.bijectors.Affine(-1.0, 2.0)
  tdist = ds.TransformedDistribution(distribution=dist, bijector=bijector)
  return tdist
//  [python]
def _build_regressor(x, z_dim=64):
  """Make the GILBO regressor, which is based off of the GAN discriminator."""
  net = tf.cast(x, tf.float32)
  net = layers.conv2d(net, 64, 4, 2, activation=tf.nn.leaky_relu)
  net = layers.conv2d(net, 128, 4, 2, activation=tf.nn.leaky_relu)
  net = layers.flatten(net)
  net = layers.dense(net, 1024, activation=tf.nn.leaky_relu)
  net = layers.dense(net, 2 * z_dim)
  # a and b correspond to the alpha beta parameters of the Beta distribution.
  a, b = net[..., :z_dim], net[..., z_dim:2 * z_dim]
  a = 1 + tf.nn.softplus(a - 5)
  b = 1 + tf.nn.softplus(b - 5)
  dist = ds.Independent(ds.Beta(a, b), 1)
  bijector = ds.bijectors.Affine(-1.0, 2.0)
  tdist = ds.TransformedDistribution(distribution=dist, bijector=bijector)
  return tdist
//  [python]
def call(self, observations, step_type, network_state):
    del step_type

    states = t.cast(t.nest.flatten(observations)[0], t.float32)
    for layer in self._dummy_layers:
      states = layer(states)

    single_action_spec = t.nest.flatten(self._output_tensor_spec)[0]
    # action_spec is TensorSpec([1], ...) so make sure there's an outer dim.
    actions = states[..., 0]
    stdevs = states[..., 1]
    actions = t.reshape(actions, [-1] + single_action_spec.shape.as_list())
    stdevs = t.reshape(stdevs, [-1] + single_action_spec.shape.as_list())
    actions = t.nest.pack_sequence_as(self._output_tensor_spec, [actions])
    stdevs = t.nest.pack_sequence_as(self._output_tensor_spec, [stdevs])

    distribution = nest_utils.map_structure_up_to(
        self._output_tensor_spec,
        tfp.distributions.MultivariateNormalDiag,
        actions,
        stdevs,
    )
    return distribution, network_state
//  [python]
def call(self, observations, step_type, network_state):
    del step_type

    states = tf.cast(tf.nest.flatten(observations)[0], tf.float32)
    for layer in self._dummy_layers:
      states = layer(states)

    single_action_spec = tf.nest.flatten(self._output_tensor_spec)[0]
    # action_spec is TensorSpec([1], ...) so make sure there's an outer dim.
    actions = states[..., 0]
    stdevs = states[..., 1]
    actions = tf.reshape(actions, [-1] + single_action_spec.shape.as_list())
    stdevs = tf.reshape(stdevs, [-1] + single_action_spec.shape.as_list())
    actions = tf.nest.pack_sequence_as(self._output_tensor_spec, [actions])
    stdevs = tf.nest.pack_sequence_as(self._output_tensor_spec, [stdevs])

    distribution = nest_utils.map_structure_up_to(
        self._output_tensor_spec,
        tfp.distributions.MultivariateNormalDiag,
        actions,
        stdevs,
    )
    return distribution, network_state
//  [python]
def correlate(x1, x2, mode="valid"):
    x1 = convert_to_tensor(x1)
    x2 = convert_to_tensor(x2)

    dtype = dtypes.result_type(
        getattr(x1, "dtype", type(x1)),
        getattr(x2, "dtype", type(x2)),
    )
    if dtype == tf.int64:
        dtype = tf.float64
    elif dtype not in [tf.bfloat16, tf.float16, tf.float64]:
        dtype = tf.float32

    x1 = tf.to_float(x1)
    x2 = tf.to_float(x2)

    x1_len, x2_len = int(x1.shape[0]), int(x2.shape[0])

    if mode == "full":
        full_len = x1_len + x2_len - 1

        x1_pad = (full_len - x1_len) / 2
        x2_pad = (full_len - x2_len) / 2

        x1 = tf.pad(
            x1, paddings=[[tf.math.floor(x1_pad), tf.math.ceil(x1_pad)]]
        )
        x2 = tf.pad(
            x2, paddings=[[tf.math.floor(x2_pad), tf.math.ceil(x2_pad)]]
        )

        x1 = tf.reshape(x1, (1, full_len, 1))
        x2 = tf.reshape(x2, (full_len, 1, 1))

        return tf.squeeze(tf.nn.conv1d(x1, x2, stride=1, padding="SAME"))

    x1 = tf.reshape(x1, (1, x1_len, 1))
    x2 = tf.reshape(x2, (x2_len, 1, 1))

    return tf.squeeze(tf.nn.conv1d(x1, x2, stride=1, padding=mode.upper()))
//  [python]
def correlate(x1, x2, mode="valid"):
    x1 = convert_to_tensor(x1)
    x2 = convert_to_tensor(x2)

    dtype = dtypes.result_type(
        getattr(x1, "dtype", type(x1)),
        getattr(x2, "dtype", type(x2)),
    )
    if dtype == tf.int64:
        dtype = tf.float64
    elif dtype not in [tf.bfloat16, tf.float16, tf.float64]:
        dtype = tf.float32

    x1 = tf.cast(x1, dtype)
    x2 = tf.cast(x2, dtype)

    x1_len, x2_len = int(x1.shape[0]), int(x2.shape[0])

    if mode == "full":
        full_len = x1_len + x2_len - 1

        x1_pad = (full_len - x1_len) / 2
        x2_pad = (full_len - x2_len) / 2

        x1 = tf.pad(
            x1, paddings=[[tf.math.floor(x1_pad), tf.math.ceil(x1_pad)]]
        )
        x2 = tf.pad(
            x2, paddings=[[tf.math.floor(x2_pad), tf.math.ceil(x2_pad)]]
        )

        x1 = tf.reshape(x1, (1, full_len, 1))
        x2 = tf.reshape(x2, (full_len, 1, 1))

        return tf.squeeze(tf.nn.conv1d(x1, x2, stride=1, padding="SAME"))

    x1 = tf.reshape(x1, (1, x1_len, 1))
    x2 = tf.reshape(x2, (x2_len, 1, 1))

    return tf.squeeze(tf.nn.conv1d(x1, x2, stride=1, padding=mode.upper()))
//  [python]
def call(
        self, inputs: TensorLike, state: TensorLike
    ) -> Tuple[tf.Tensor, tf.Tensor]:
        """Build the CrfDecodeForwardRnnCell.

        Args:
          inputs: A [batch_size, num_tags] matrix of unary potentials.
          state: A [batch_size, num_tags] matrix containing the previous step's
                score values.

        Returns:
          output: A [batch_size, num_tags * 2] matrix of backpointers and scores.
          new_state: A [batch_size, num_tags] matrix of new score values.
        """
        state = tf.expand_dims(state[0], 2)
        transition_scores = state + self._transition_params
        new_state = inputs + tf.reduce_max(transition_scores, [1])

        backpointers = tf.argmax(transition_scores, 1)
        backpointers = tf.cast(backpointers, tf.float32)

        # apply softmax to transition_scores to get scores in range from 0 to 1
        scores = tf.reduce_max(tf.nn.softmax(transition_scores, axis=1), [1])

        # In the RNN implementation only the first value that is returned from a cell
        # is kept throughout the RNN, so that you will have the values from each time
        # step in the final output. As we need the backpointers as well as the scores
        # for each time step, we concatenate them.
        return tf.concat([backpointers, scores], axis=1), new_state
//  [python]
def crf_decode_forward(
    inputs: TensorLike,
    state: TensorLike,
    transition_params: TensorLike,
    sequence_lengths: TensorLike,
) -> Tuple[tf.Tensor, tf.Tensor]:
    """Computes forward decoding in a linear-chain CRF.

    Args:
      inputs: A [batch_size, num_tags] matrix of unary potentials.
      state: A [batch_size, num_tags] matrix containing the previous step's
            score values.
      transition_params: A [num_tags, num_tags] matrix of binary potentials.
      sequence_lengths: A [batch_size] vector of true sequence lengths.

    Returns:
      output: A [batch_size, num_tags * 2] matrix of backpointers and scores.
      new_state: A [batch_size, num_tags] matrix of new score values.
    """
    sequence_lengths = t.cast(sequence_lengths, dtype=t.int32)
    mask = tf.sequence_mask(sequence_lengths, tf.shape(inputs)[1])
    crf_fwd_cell = CrfDecodeForwardRnnCell(transition_params)
    crf_fwd_layer = tf.keras.layers.RNN(
        crf_fwd_cell, return_sequences=True, return_state=True
    )
    return crf_fwd_layer(inputs, state, mask=mask)
//  [python]
def crf_decode_forward(
    inputs: TensorLike,
    state: TensorLike,
    transition_params: TensorLike,
    sequence_lengths: TensorLike,
) -> Tuple[tf.Tensor, tf.Tensor]:
    """Computes forward decoding in a linear-chain CRF.

    Args:
      inputs: A [batch_size, num_tags] matrix of unary potentials.
      state: A [batch_size, num_tags] matrix containing the previous step's
            score values.
      transition_params: A [num_tags, num_tags] matrix of binary potentials.
      sequence_lengths: A [batch_size] vector of true sequence lengths.

    Returns:
      output: A [batch_size, num_tags * 2] matrix of backpointers and scores.
      new_state: A [batch_size, num_tags] matrix of new score values.
    """
    sequence_lengths = tf.cast(sequence_lengths, dtype=tf.int32)
    mask = tf.sequence_mask(sequence_lengths, tf.shape(inputs)[1])
    crf_fwd_cell = CrfDecodeForwardRnnCell(transition_params)
    crf_fwd_layer = tf.keras.layers.RNN(
        crf_fwd_cell, return_sequences=True, return_state=True
    )
    return crf_fwd_layer(inputs, state, mask=mask)
//  [python]
def crf_decode_backward(
    backpointers: TensorLike, scores: TensorLike, state: TensorLike
) -> Tuple[tf.Tensor, tf.Tensor]:
    """Computes backward decoding in a linear-chain CRF.

    Args:
      backpointers: A [batch_size, num_tags] matrix of backpointer of next step
            (in time order).
      scores: A [batch_size, num_tags] matrix of scores of next step (in time order).
      state: A [batch_size, 1] matrix of tag index of next step.

    Returns:
      new_tags: A [batch_size, num_tags] tensor containing the new tag indices.
      new_scores: A [batch_size, num_tags] tensor containing the new score values.
    """
    backpointers = tf.transpose(backpointers, [1, 0, 2])
    scores = tf.transpose(scores, [1, 0, 2])

    def _scan_fn(_state: TensorLike, _inputs: TensorLike) -> tf.Tensor:
        _state = tf.squeeze(_state, axis=[1])
        idxs = tf.stack([tf.range(tf.shape(_inputs)[0]), _state], axis=1)
        return tf.expand_dims(tf.gather_nd(_inputs, idxs), axis=-1)

    output_tags = tf.scan(_scan_fn, backpointers, state)
    # the dtype of the input parameters of tf.scan need to match
    # convert state to float32 to match the type of scores
    state = tf.to_float(state)
    output_scores = tf.scan(_scan_fn, scores, state)

    return tf.transpose(output_tags, [1, 0, 2]), tf.transpose(output_scores, [1, 0, 2])
//  [python]
def crf_decode_backward(
    backpointers: TensorLike, scores: TensorLike, state: TensorLike
) -> Tuple[tf.Tensor, tf.Tensor]:
    """Computes backward decoding in a linear-chain CRF.

    Args:
      backpointers: A [batch_size, num_tags] matrix of backpointer of next step
            (in time order).
      scores: A [batch_size, num_tags] matrix of scores of next step (in time order).
      state: A [batch_size, 1] matrix of tag index of next step.

    Returns:
      new_tags: A [batch_size, num_tags] tensor containing the new tag indices.
      new_scores: A [batch_size, num_tags] tensor containing the new score values.
    """
    backpointers = tf.transpose(backpointers, [1, 0, 2])
    scores = tf.transpose(scores, [1, 0, 2])

    def _scan_fn(_state: TensorLike, _inputs: TensorLike) -> tf.Tensor:
        _state = tf.cast(tf.squeeze(_state, axis=[1]), dtype=tf.int32)
        idxs = tf.stack([tf.range(tf.shape(_inputs)[0]), _state], axis=1)
        return tf.expand_dims(tf.gather_nd(_inputs, idxs), axis=-1)

    output_tags = tf.scan(_scan_fn, backpointers, state)
    # the dtype of the input parameters of tf.scan need to match
    # convert state to float32 to match the type of scores
    state = tf.cast(state, dtype=tf.float32)
    output_scores = tf.scan(_scan_fn, scores, state)

    return tf.transpose(output_tags, [1, 0, 2]), tf.transpose(output_scores, [1, 0, 2])
//  [python]
def policy_frac(hparams):
    """How far we are through policy training."""
    return tf.cast(tf.train.get_global_step(), tf.float32) / nupdates(hparams)
//  [python]
def crf_binary_score(
    tag_indices: TensorLike, sequence_lengths: TensorLike, transition_params: TensorLike
) -> tf.Tensor:
    """Computes the binary scores of tag sequences.

    Args:
      tag_indices: A [batch_size, max_seq_len] matrix of tag indices.
      sequence_lengths: A [batch_size] vector of true sequence lengths.
      transition_params: A [num_tags, num_tags] matrix of binary potentials.

    Returns:
      binary_scores: A [batch_size] vector of binary scores.
    """
    tag_indices = t.cast(tag_indices, dtype=tf.int32)
    sequence_lengths = t.cast(sequence_lengths, dtype=tf.int32)

    num_tags = tf.shape(transition_params)[0]
    num_transitions = tf.shape(tag_indices)[1] - 1

    # Truncate by one on each side of the sequence to get the start and end
    # indices of each transition.
    start_tag_indices = tf.slice(tag_indices, [0, 0], [-1, num_transitions])
    end_tag_indices = tf.slice(tag_indices, [0, 1], [-1, num_transitions])

    # Encode the indices in a flattened representation.
    flattened_transition_indices = start_tag_indices * num_tags + end_tag_indices
    flattened_transition_params = tf.reshape(transition_params, [-1])

    # Get the binary scores based on the flattened representation.
    binary_scores = tf.gather(flattened_transition_params, flattened_transition_indices)

    masks = tf.sequence_mask(
        sequence_lengths, maxlen=tf.shape(tag_indices)[1], dtype=binary_scores.dtype
    )
    truncated_masks = tf.slice(masks, [0, 1], [-1, -1])
    binary_scores = tf.reduce_sum(binary_scores * truncated_masks, 1)
    return binary_scores
//  [python]
def crf_binary_score(
    tag_indices: TensorLike, sequence_lengths: TensorLike, transition_params: TensorLike
) -> tf.Tensor:
    """Computes the binary scores of tag sequences.

    Args:
      tag_indices: A [batch_size, max_seq_len] matrix of tag indices.
      sequence_lengths: A [batch_size] vector of true sequence lengths.
      transition_params: A [num_tags, num_tags] matrix of binary potentials.

    Returns:
      binary_scores: A [batch_size] vector of binary scores.
    """
    tag_indices = tf.cast(tag_indices, dtype=tf.int32)
    sequence_lengths = tf.cast(sequence_lengths, dtype=tf.int32)

    num_tags = tf.shape(transition_params)[0]
    num_transitions = tf.shape(tag_indices)[1] - 1

    # Truncate by one on each side of the sequence to get the start and end
    # indices of each transition.
    start_tag_indices = tf.slice(tag_indices, [0, 0], [-1, num_transitions])
    end_tag_indices = tf.slice(tag_indices, [0, 1], [-1, num_transitions])

    # Encode the indices in a flattened representation.
    flattened_transition_indices = start_tag_indices * num_tags + end_tag_indices
    flattened_transition_params = tf.reshape(transition_params, [-1])

    # Get the binary scores based on the flattened representation.
    binary_scores = tf.gather(flattened_transition_params, flattened_transition_indices)

    masks = tf.sequence_mask(
        sequence_lengths, maxlen=tf.shape(tag_indices)[1], dtype=binary_scores.dtype
    )
    truncated_masks = tf.slice(masks, [0, 1], [-1, -1])
    binary_scores = tf.reduce_sum(binary_scores * truncated_masks, 1)
    return binary_scores
//  [python]
def crf_unary_score(
    tag_indices: TensorLike, sequence_lengths: TensorLike, inputs: TensorLike
) -> tf.Tensor:
    """Computes the unary scores of tag sequences.

    Args:
      tag_indices: A [batch_size, max_seq_len] matrix of tag indices.
      sequence_lengths: A [batch_size] vector of true sequence lengths.
      inputs: A [batch_size, max_seq_len, num_tags] tensor of unary potentials.

    Returns:
      unary_scores: A [batch_size] vector of unary scores.
    """
    tag_indices = t.cast(tag_indices, dtype=tf.int32)
    sequence_lengths = t.cast(sequence_lengths, dtype=tf.int32)

    batch_size = tf.shape(inputs)[0]
    max_seq_len = tf.shape(inputs)[1]
    num_tags = tf.shape(inputs)[2]

    flattened_inputs = tf.reshape(inputs, [-1])

    offsets = tf.expand_dims(tf.range(batch_size) * max_seq_len * num_tags, 1)
    offsets += tf.expand_dims(tf.range(max_seq_len) * num_tags, 0)
    # Use int32 or int64 based on tag_indices' dtype.
    if tag_indices.dtype == tf.int64:
        offsets = t.cast(offsets, tf.int64)
    flattened_tag_indices = tf.reshape(offsets + tag_indices, [-1])

    unary_scores = tf.reshape(
        tf.gather(flattened_inputs, flattened_tag_indices), [batch_size, max_seq_len]
    )

    masks = tf.sequence_mask(
        sequence_lengths, maxlen=tf.shape(tag_indices)[1], dtype=unary_scores.dtype
    )

    unary_scores = tf.reduce_sum(unary_scores * masks, 1)
    return unary_scores
//  [python]
def crf_unary_score(
    tag_indices: TensorLike, sequence_lengths: TensorLike, inputs: TensorLike
) -> tf.Tensor:
    """Computes the unary scores of tag sequences.

    Args:
      tag_indices: A [batch_size, max_seq_len] matrix of tag indices.
      sequence_lengths: A [batch_size] vector of true sequence lengths.
      inputs: A [batch_size, max_seq_len, num_tags] tensor of unary potentials.

    Returns:
      unary_scores: A [batch_size] vector of unary scores.
    """
    tag_indices = tf.cast(tag_indices, dtype=tf.int32)
    sequence_lengths = tf.cast(sequence_lengths, dtype=tf.int32)

    batch_size = tf.shape(inputs)[0]
    max_seq_len = tf.shape(inputs)[1]
    num_tags = tf.shape(inputs)[2]

    flattened_inputs = tf.reshape(inputs, [-1])

    offsets = tf.expand_dims(tf.range(batch_size) * max_seq_len * num_tags, 1)
    offsets += tf.expand_dims(tf.range(max_seq_len) * num_tags, 0)
    # Use int32 or int64 based on tag_indices' dtype.
    if tag_indices.dtype == tf.int64:
        offsets = tf.cast(offsets, tf.int64)
    flattened_tag_indices = tf.reshape(offsets + tag_indices, [-1])

    unary_scores = tf.reshape(
        tf.gather(flattened_inputs, flattened_tag_indices), [batch_size, max_seq_len]
    )

    masks = tf.sequence_mask(
        sequence_lengths, maxlen=tf.shape(tag_indices)[1], dtype=unary_scores.dtype
    )

    unary_scores = tf.reduce_sum(unary_scores * masks, 1)
    return unary_scores
//  [python]
def crf_forward(
    inputs: TensorLike,
    state: TensorLike,
    transition_params: TensorLike,
    sequence_lengths: TensorLike,
) -> tf.Tensor:
    """Computes the alpha values in a linear-chain CRF.

    See http://www.cs.columbia.edu/~mcollins/fb.pdf for reference.

    Args:
      inputs: A [batch_size, num_tags] matrix of unary potentials.
      state: A [batch_size, num_tags] matrix containing the previous alpha
         values.
      transition_params: A [num_tags, num_tags] matrix of binary potentials.
          This matrix is expanded into a [1, num_tags, num_tags] in preparation
          for the broadcast summation occurring within the cell.
      sequence_lengths: A [batch_size] vector of true sequence lengths.

    Returns:
      new_alphas: A [batch_size, num_tags] matrix containing the
          new alpha values.
    """
    sequence_lengths = t.cast(sequence_lengths, dtype=t.int32)

    last_index = tf.maximum(
        tf.constant(0, dtype=sequence_lengths.dtype), sequence_lengths - 1
    )
    inputs = tf.transpose(inputs, [1, 0, 2])
    transition_params = tf.expand_dims(transition_params, 0)

    def _scan_fn(_state: TensorLike, _inputs: TensorLike) -> TensorLike:
        _state = tf.expand_dims(_state, 2)
        transition_scores = _state + transition_params
        new_alphas = _inputs + tf.reduce_logsumexp(transition_scores, [1])
        return new_alphas

    all_alphas = tf.transpose(tf.scan(_scan_fn, inputs, state), [1, 0, 2])
    # add first state for sequences of length 1
    all_alphas = tf.concat([tf.expand_dims(state, 1), all_alphas], 1)

    idxs = tf.stack([tf.range(tf.shape(last_index)[0]), last_index], axis=1)
    return tf.gather_nd(all_alphas, idxs)
//  [python]
def crf_forward(
    inputs: TensorLike,
    state: TensorLike,
    transition_params: TensorLike,
    sequence_lengths: TensorLike,
) -> tf.Tensor:
    """Computes the alpha values in a linear-chain CRF.

    See http://www.cs.columbia.edu/~mcollins/fb.pdf for reference.

    Args:
      inputs: A [batch_size, num_tags] matrix of unary potentials.
      state: A [batch_size, num_tags] matrix containing the previous alpha
         values.
      transition_params: A [num_tags, num_tags] matrix of binary potentials.
          This matrix is expanded into a [1, num_tags, num_tags] in preparation
          for the broadcast summation occurring within the cell.
      sequence_lengths: A [batch_size] vector of true sequence lengths.

    Returns:
      new_alphas: A [batch_size, num_tags] matrix containing the
          new alpha values.
    """
    sequence_lengths = tf.cast(sequence_lengths, dtype=tf.int32)

    last_index = tf.maximum(
        tf.constant(0, dtype=sequence_lengths.dtype), sequence_lengths - 1
    )
    inputs = tf.transpose(inputs, [1, 0, 2])
    transition_params = tf.expand_dims(transition_params, 0)

    def _scan_fn(_state: TensorLike, _inputs: TensorLike) -> TensorLike:
        _state = tf.expand_dims(_state, 2)
        transition_scores = _state + transition_params
        new_alphas = _inputs + tf.reduce_logsumexp(transition_scores, [1])
        return new_alphas

    all_alphas = tf.transpose(tf.scan(_scan_fn, inputs, state), [1, 0, 2])
    # add first state for sequences of length 1
    all_alphas = tf.concat([tf.expand_dims(state, 1), all_alphas], 1)

    idxs = tf.stack([tf.range(tf.shape(last_index)[0]), last_index], axis=1)
    return tf.gather_nd(all_alphas, idxs)
//  [python]
def crf_log_likelihood(
    inputs: TensorLike,
    tag_indices: TensorLike,
    sequence_lengths: TensorLike,
    transition_params: Optional[TensorLike] = None,
) -> Tuple[tf.Tensor, tf.Tensor]:
    """Computes the log-likelihood of tag sequences in a CRF.

    Args:
      inputs: A [batch_size, max_seq_len, num_tags] tensor of unary potentials
          to use as input to the CRF layer.
      tag_indices: A [batch_size, max_seq_len] matrix of tag indices for which
          we compute the log-likelihood.
      sequence_lengths: A [batch_size] vector of true sequence lengths.
      transition_params: A [num_tags, num_tags] transition matrix,
          if available.

    Returns:
      log_likelihood: A [batch_size] `Tensor` containing the log-likelihood of
        each example, given the sequence of tag indices.
      transition_params: A [num_tags, num_tags] transition matrix. This is
          either provided by the caller or created in this function.
    """
    inputs = tf.convert_to_tensor(inputs)

    num_tags = inputs.shape[2]

    # cast type to handle different types
    tag_indices = tf.to_int32(tag_indices)
    sequence_lengths = tf.to_int32(sequence_lengths)

    if transition_params is None:
        initializer = tf.keras.initializers.GlorotUniform()
        transition_params = tf.Variable(
            initializer([num_tags, num_tags]), "transitions"
        )
    transition_params = tf.cast(transition_params, inputs.dtype)
    sequence_scores = crf_sequence_score(
        inputs, tag_indices, sequence_lengths, transition_params
    )
    log_norm = crf_log_norm(inputs, sequence_lengths, transition_params)

    # Normalize the scores to get the log-likelihood per example.
    log_likelihood = sequence_scores - log_norm
    return log_likelihood, transition_params
//  [python]
def crf_log_likelihood(
    inputs: TensorLike,
    tag_indices: TensorLike,
    sequence_lengths: TensorLike,
    transition_params: Optional[TensorLike] = None,
) -> Tuple[tf.Tensor, tf.Tensor]:
    """Computes the log-likelihood of tag sequences in a CRF.

    Args:
      inputs: A [batch_size, max_seq_len, num_tags] tensor of unary potentials
          to use as input to the CRF layer.
      tag_indices: A [batch_size, max_seq_len] matrix of tag indices for which
          we compute the log-likelihood.
      sequence_lengths: A [batch_size] vector of true sequence lengths.
      transition_params: A [num_tags, num_tags] transition matrix,
          if available.

    Returns:
      log_likelihood: A [batch_size] `Tensor` containing the log-likelihood of
        each example, given the sequence of tag indices.
      transition_params: A [num_tags, num_tags] transition matrix. This is
          either provided by the caller or created in this function.
    """
    inputs = tf.convert_to_tensor(inputs)

    num_tags = inputs.shape[2]

    # cast type to handle different types
    tag_indices = tf.cast(tag_indices, dtype=tf.int32)
    sequence_lengths = tf.cast(sequence_lengths, dtype=tf.int32)

    if transition_params is None:
        initializer = tf.keras.initializers.GlorotUniform()
        transition_params = tf.Variable(
            initializer([num_tags, num_tags]), "transitions"
        )
    transition_params = tf.cast(transition_params, inputs.dtype)
    sequence_scores = crf_sequence_score(
        inputs, tag_indices, sequence_lengths, transition_params
    )
    log_norm = crf_log_norm(inputs, sequence_lengths, transition_params)

    # Normalize the scores to get the log-likelihood per example.
    log_likelihood = sequence_scores - log_norm
    return log_likelihood, transition_params
//  [python]
def crf_sequence_score(
    inputs: TensorLike,
    tag_indices: TensorLike,
    sequence_lengths: TensorLike,
    transition_params: TensorLike,
) -> tf.Tensor:
    """Computes the unnormalized score for a tag sequence.

    Args:
      inputs: A [batch_size, max_seq_len, num_tags] tensor of unary potentials
          to use as input to the CRF layer.
      tag_indices: A [batch_size, max_seq_len] matrix of tag indices for which
          we compute the unnormalized score.
      sequence_lengths: A [batch_size] vector of true sequence lengths.
      transition_params: A [num_tags, num_tags] transition matrix.

    Returns:
      sequence_scores: A [batch_size] vector of unnormalized sequence scores.
    """
    tag_indices = tf.to_int32(tag_indices)
    sequence_lengths = tf.to_int32(sequence_lengths)

    # If max_seq_len is 1, we skip the score calculation and simply gather the
    # unary potentials of the single tag.
    def _single_seq_fn() -> TensorLike:
        batch_size = tf.shape(inputs, out_type=tf.int32)[0]
        batch_inds = tf.reshape(tf.range(batch_size), [-1, 1])
        indices = tf.concat([batch_inds, tf.zeros_like(batch_inds)], axis=1)

        tag_inds = tf.gather_nd(tag_indices, indices)
        tag_inds = tf.reshape(tag_inds, [-1, 1])
        indices = tf.concat([indices, tag_inds], axis=1)

        sequence_scores = tf.gather_nd(inputs, indices)

        sequence_scores = tf.where(
            tf.less_equal(sequence_lengths, 0),
            tf.zeros_like(sequence_scores),
            sequence_scores,
        )
        return sequence_scores

    def _multi_seq_fn() -> TensorLike:
        # Compute the scores of the given tag sequence.
        unary_scores = crf_unary_score(tag_indices, sequence_lengths, inputs)
        binary_scores = crf_binary_score(
            tag_indices, sequence_lengths, transition_params
        )
        sequence_scores = unary_scores + binary_scores
        return sequence_scores

    return tf.cond(tf.equal(tf.shape(inputs)[1], 1), _single_seq_fn, _multi_seq_fn)