Core ops

keras.ops.cast(x, dtype)

Cast a tensor to the desired dtype.

Arguments

• x: A tensor or variable.
• dtype: The target type.

Returns

A tensor of the specified dtype.

Example

>>> x = keras.ops.arange(4)
>>> x = keras.ops.cast(x, dtype="float16")

cond function

keras.ops.cond(pred, true_fn, false_fn)

Conditionally applies true_fn or false_fn.

Arguments

• pred: Boolean scalar type
• true_fn: Callable returning the output for the pred == True case.
• false_fn: Callable returning the output for the pred == False case.

Returns

The output of either true_fn or false_fn depending on pred.

convert_to_numpy function

keras.ops.convert_to_numpy(x)

Convert a tensor to a NumPy array.

Arguments

• x: A tensor.

Returns

A NumPy array.

convert_to_tensor function

keras.ops.convert_to_tensor(x, dtype=None, sparse=None)

Convert a NumPy array to a tensor.

Arguments

• x: A NumPy array.
• dtype: The target type.
• sparse: Whether to keep sparse tensors. False will cause sparse tensors to be densified. The default value of None means that sparse tensors are kept only if the backend supports them.

Returns

A tensor of the specified dtype.

Example

>>> x = np.array([1, 2, 3])
>>> y = keras.ops.convert_to_tensor(x)

erf function

keras.ops.erf(x)

Computes the error function of x, element-wise.

Arguments

• x: Input tensor.

Returns

A tensor with the same dtype as x.

Example

>>> x = np.array([-3.0, -2.0, -1.0, 0.0, 1.0])
>>> keras.ops.erf(x)
array([-0.99998 , -0.99532, -0.842701,  0.,  0.842701], dtype=float32)

erfinv function

keras.ops.erfinv(x)

Computes the inverse error function of x, element-wise.

Arguments

• x: Input tensor.

Returns

A tensor with the same dtype as x.

Example

>>> x = np.array([-0.5, -0.2, -0.1, 0.0, 0.3])
>>> keras.ops.erfinv(x)
array([-0.47694, -0.17914, -0.08886,  0. ,  0.27246], dtype=float32)

extract_sequences function

keras.ops.extract_sequences(x, sequence_length, sequence_stride)

Expands the dimension of last axis into sequences of sequence_length.

Slides a window of size sequence_length over the last axis of the input with a stride of sequence_stride, replacing the last axis with [num_sequences, sequence_length] sequences.

If the dimension along the last axis is N, the number of sequences can be computed by:

num_sequences = 1 + (N - sequence_length) // sequence_stride

Arguments

• x: Input tensor.
• sequence_length: An integer representing the sequences length.
• sequence_stride: An integer representing the sequences hop size.

Returns

A tensor of sequences with shape […, num_sequences, sequence_length].

Example

>>> x = keras.ops.convert_to_tensor([1, 2, 3, 4, 5, 6])
>>> extract_sequences(x, 3, 2)
array([[1, 2, 3],
   [3, 4, 5]])

fori_loop function

keras.ops.fori_loop(lower, upper, body_fun, init_val)

For loop implementation.

Arguments

• lower: The initial value of the loop variable.
• upper: The upper bound of the loop variable.
• body_fun: A callable that represents the loop body. Must take two arguments: the loop variable and the loop state. The loop state should be updated and returned by this function.
• init_val: The initial value of the loop state.

Returns

The final state after the loop.

Example

>>> lower = 0
>>> upper = 10
>>> body_fun = lambda i, s: (i + 1, s + i)
>>> init_val = 0
>>> keras.ops.fori_loop(lower, upper, body_fun, init_val)
45

in_top_k function

keras.ops.in_top_k(targets, predictions, k)

Checks if the targets are in the top-k predictions.

Arguments

• targets: A tensor of true labels.
• predictions: A tensor of predicted labels.
• k: An integer representing the number of predictions to consider.

Returns

A boolean tensor of the same shape as targets, where each element indicates whether the corresponding target is in the top-k predictions.

Example

>>> targets = keras.ops.convert_to_tensor([2, 5, 3])
>>> predictions = keras.ops.convert_to_tensor(
... [[0.1, 0.4, 0.6, 0.9, 0.5],
...  [0.1, 0.7, 0.9, 0.8, 0.3],
...  [0.1, 0.6, 0.9, 0.9, 0.5]])
>>> in_top_k(targets, predictions, k=3)
array([ True False  True], shape=(3,), dtype=bool)

is_tensor function

keras.ops.is_tensor(x)

Check whether the given object is a tensor.

Note: This checks for backend specific tensors so passing a TensorFlow tensor would return False if your backend is PyTorch or JAX.

Arguments

• x: A variable.

Returns

True if x is a tensor, otherwise False.

logsumexp function

keras.ops.logsumexp(x, axis=None, keepdims=False)

Computes the logarithm of sum of exponentials of elements in a tensor.

Arguments

• x: Input tensor.
• axis: An integer or a tuple of integers specifying the axis/axes along which to compute the sum. If None, the sum is computed over all elements. Defaults toNone.
• keepdims: A boolean indicating whether to keep the dimensions of the input tensor when computing the sum. Defaults toFalse.

Returns

A tensor containing the logarithm of the sum of exponentials of elements in x.

Example

>>> x = keras.ops.convert_to_tensor([1., 2., 3.])
>>> logsumexp(x)
3.407606

rsqrt function

keras.ops.rsqrt(x)

Computes reciprocal of square root of x element-wise.

Arguments

• x: input tensor

Returns

A tensor with the same dtype as x.

Example

>>> x = keras.ops.convert_to_tensor([1.0, 10.0, 100.0])
>>> keras.ops.rsqrt(x)
array([1.0, 0.31622776, 0.1], dtype=float32)

scatter function

keras.ops.scatter(indices, values, shape)

Returns a tensor of shape shape where indices are set to values.

At a high level, this operation does zeros[indices] = updates and returns the output. It is equivalent to:

zeros = keras.ops.zeros(shape)
output = keras.ops.scatter_update(zeros, indices, values)

Arguments

• indices: A tensor or list/tuple specifying indices for the values in values.
• values: A tensor, the values to be set at indices.
• shape: Shape of the output tensor.

Example

>>> indices = [[0, 1], [1, 1]]
>>> values = np.array([1., 1.])
>>> keras.ops.scatter(indices, values, shape=(2, 2))
array([[0., 1.],
       [0., 1.]])

scatter_update function

keras.ops.scatter_update(inputs, indices, updates)

Update inputs via updates at scattered (sparse) indices.

At a high level, this operation does inputs[indices] = updates. Assume inputs is a tensor of shape (D0, D1, ..., Dn), there are 2 main usages of scatter_update.

1. indices is a 2D tensor of shape (num_updates, n), where num_updates is the number of updates to perform, and updates is a 1D tensor of shape (num_updates,). For example, if inputs is zeros((4, 4, 4)), and we want to update inputs[1, 2, 3] and inputs[0, 1, 3] as 1, then we can use:

inputs = np.zeros((4, 4, 4))
indices = [[1, 2, 3], [0, 1, 3]]
updates = np.array([1., 1.])
inputs = keras.ops.scatter_update(inputs, indices, updates)

2 indices is a 2D tensor of shape (num_updates, k), where num_updates is the number of updates to perform, and k (k < n) is the size of each index in indices. updates is a n - k-D tensor of shape (num_updates, inputs.shape[k:]). For example, if inputs = np.zeros((4, 4, 4)), and we want to update inputs[1, 2, :] and inputs[2, 3, :] as [1, 1, 1, 1], then indices would have shape (num_updates, 2) (k = 2), and updates would have shape (num_updates, 4) (inputs.shape[2:] = 4). See the code below:

inputs = np.zeros((4, 4, 4))
indices = [[1, 2], [2, 3]]
updates = np.array([[1., 1., 1, 1,], [1., 1., 1, 1,])
inputs = keras.ops.scatter_update(inputs, indices, updates)

Arguments

• inputs: A tensor, the tensor to be updated.
• indices: A tensor or list/tuple of shape (N, inputs.ndim), specifying indices to update. N is the number of indices to update, must be equal to the first dimension of updates.
• updates: A tensor, the new values to be put to inputs at indices.

Returns

A tensor, has the same shape and dtype as inputs.

segment_max function

keras.ops.segment_max(data, segment_ids, num_segments=None, sorted=False)

Computes the max of segments in a tensor.

Arguments

• data: Input tensor.
• segment_ids: A 1-D tensor containing segment indices for each element in data.
• num_segments: An integer representing the total number of segments. If not specified, it is inferred from the maximum value in segment_ids.
• sorted: A boolean indicating whether segment_ids is sorted. Defaults toFalse.

Returns

A tensor containing the max of segments, where each element represents the max of the corresponding segment in data.

Example

>>> data = keras.ops.convert_to_tensor([1, 2, 10, 20, 100, 200])
>>> segment_ids = keras.ops.convert_to_tensor([0, 0, 1, 1, 2, 2])
>>> num_segments = 3
>>> keras.ops.segment_max(data, segment_ids, num_segments)
array([2, 20, 200], dtype=int32)

segment_sum function

keras.ops.segment_sum(data, segment_ids, num_segments=None, sorted=False)

Computes the sum of segments in a tensor.

Arguments

• data: Input tensor.
• segment_ids: A 1-D tensor containing segment indices for each element in data.
• num_segments: An integer representing the total number of segments. If not specified, it is inferred from the maximum value in segment_ids.
• sorted: A boolean indicating whether segment_ids is sorted. Defaults toFalse.

Returns

A tensor containing the sum of segments, where each element represents the sum of the corresponding segment in data.

Example

>>> data = keras.ops.convert_to_tensor([1, 2, 10, 20, 100, 200])
>>> segment_ids = keras.ops.convert_to_tensor([0, 0, 1, 1, 2, 2])
>>> num_segments = 3
>>> keras.ops.segment_sum(data, segment_ids,num_segments)
array([3, 30, 300], dtype=int32)

shape function

keras.ops.shape(x)

Gets the shape of the tensor input.

Note: On the TensorFlow backend, when x is a tf.Tensor with dynamic shape, dimensions which are dynamic in the context of a compiled function will have a tf.Tensor value instead of a static integer value.

Arguments

• x: A tensor. This function will try to access the shape attribute of the input tensor.

Returns

A tuple of integers or None values, indicating the shape of the input tensor.

Example

>>> x = keras.zeros((8, 12))
>>> keras.ops.shape(x)
(8, 12)

slice function

keras.ops.slice(inputs, start_indices, shape)

Return a slice of an input tensor.

At a high level, this operation is an explicit replacement for array slicing e.g. inputs[start_indices: start_indices + shape]. Unlike slicing via brackets, this operation will accept tensor start indices on all backends, which is useful when indices dynamically computed via other tensor operations.

inputs = np.zeros((5, 5))
start_indices = np.array([3, 3])
shape = np.array([2, 2])
inputs = keras.ops.slice(inputs, start_indices, updates)

Arguments

• inputs: A tensor, the tensor to be updated.
• start_indices: A list/tuple of shape (inputs.ndim,), specifying the starting indices for updating.
• shape: The full shape of the returned slice.

Returns

A tensor, has the same shape and dtype as inputs.

slice_update function

keras.ops.slice_update(inputs, start_indices, updates)

Update an input by slicing in a tensor of updated values.

At a high level, this operation does inputs[start_indices: start_indices + updates.shape] = updates. Assume inputs is a tensor of shape (D0, D1, ..., Dn), start_indices must be a list/tuple of n integers, specifying the starting indices. updates must have the same rank as inputs, and the size of each dim must not exceed Di - start_indices[i]. For example, if we have 2D inputs inputs = np.zeros((5, 5)), and we want to update the intersection of last 2 rows and last 2 columns as 1, i.e., inputs[3:, 3:] = np.ones((2, 2)), then we can use the code below:

inputs = np.zeros((5, 5))
start_indices = [3, 3]
updates = np.ones((2, 2))
inputs = keras.ops.slice_update(inputs, start_indices, updates)

Arguments

• inputs: A tensor, the tensor to be updated.
• start_indices: A list/tuple of shape (inputs.ndim,), specifying the starting indices for updating.
• updates: A tensor, the new values to be put to inputs at indices. updates must have the same rank as inputs.

Returns

A tensor, has the same shape and dtype as inputs.

stop_gradient function

keras.ops.stop_gradient(variable)

Stops gradient computation.

Arguments

• variable: A tensor variable for which the gradient computation is to be disabled.

Returns

The variable with gradient computation disabled.

Examples

>>> var = keras.backend.convert_to_tensor(
...     [1., 2., 3.],
...     dtype="float32"
... )
>>> var = keras.ops.stop_gradient(var)

top_k function

keras.ops.top_k(x, k, sorted=True)

Finds the top-k values and their indices in a tensor.

Arguments

• x: Input tensor.
• k: An integer representing the number of top elements to retrieve.
• sorted: A boolean indicating whether to sort the output in descending order. Defaults toTrue.

Returns

A tuple containing two tensors. The first tensor contains the top-k values, and the second tensor contains the indices of the top-k values in the input tensor.

Example

>>> x = keras.ops.convert_to_tensor([5, 2, 7, 1, 9, 3])
>>> values, indices = top_k(x, k=3)
>>> print(values)
array([9 7 5], shape=(3,), dtype=int32)
>>> print(indices)
array([4 2 0], shape=(3,), dtype=int32)

unstack function

keras.ops.unstack(x, num=None, axis=0)

Unpacks the given dimension of a rank-R tensor into rank-(R-1) tensors.

Arguments

• x: The input tensor.
• num: The length of the dimension axis. Automatically inferred if None.
• axis: The axis along which to unpack.

Returns

A list of tensors unpacked along the given axis.

Example

>>> x = keras.ops.array([[1, 2], [3, 4]])
>>> keras.ops.unstack(x, axis=0)
[array([1, 2]), array([3, 4])]

vectorized_map function

keras.ops.vectorized_map(function, elements)

Parallel map of function on axis 0 of tensor(s) elements.

Schematically, vectorized_map implements the following, in the case of a single tensor input elements:

def vectorized_map(function, elements)
    outputs = []
    for e in elements:
        outputs.append(function(e))
    return stack(outputs)

In the case of an iterable of tensors elements, it implements the following:

def vectorized_map(function, elements)
    batch_size = elements[0].shape[0]
    outputs = []
    for index in range(batch_size):
        outputs.append(function([e[index] for e in elements]))
    return np.stack(outputs)

In this case, function is expected to take as input a single list of tensor arguments.

while_loop function

keras.ops.while_loop(cond, body, loop_vars, maximum_iterations=None)

While loop implementation.

Arguments

• cond: A callable that represents the termination condition of the loop. Must accept a loop_vars like structure as an argument. If loop_vars is a tuple or list, each element of loop_vars will be passed positionally to the callable.
• body: A callable that represents the loop body. Must accept a loop_vars like structure as an argument, and return update value with the same structure. If loop_vars is a tuple or list, each element of loop_vars will be passed positionally to the callable.
• loop_vars: An arbitrary nested structure of tensor state to persist across loop iterations.
• maximum_iterations: Optional maximum number of iterations of the while loop to run. If provided, the cond output is AND-ed with an additional condition ensuring the number of iterations executed is no greater than maximum_iterations.

Returns

A list/tuple of tensors, has the same shape and dtype as inputs.

Examples

>>> i = 0
>>> cond = lambda i: i < 10
>>> body = lambda i: i + 1
>>> keras.ops.while_loop(cond, body, i)
10

>>> x, y = 0, 1
>>> cond = lambda x, y: x < 10
>>> body = lambda x, y: (x + 1, y + 1)
>>> keras.ops.while_loop(cond, body, (x, y))
10, 11