ttnn.where_bw

ttnn.where_bw(grad_tensor: ttnn.Tensor, input_tensor_a: ttnn.Tensor, input_tensor_b: ttnn.Tensor, input_tensor_c: ttnn.Tensor, *, memory_config: ttnn.MemoryConfig = None, are_required_outputs: List[bool] = [True, True], input_a_grad: ttnn.Tensor = None, input_b_grad: ttnn.Tensor = None) → List of ttnn.Tensor

Performs backward operations for where of input_tensor_a, input_tensor_b and input_tensor_c with given grad_tensor.

The three input tensors are the arguments of a forward ttnn.where call, with input_tensor_a as its predicate, and grad_tensor is the gradient with respect to the output of that call, so its elements line up one to one with that output. The forward ttnn.where copies each output element from input_tensor_b where the predicate is 1 and from input_tensor_c where the predicate is 0, so each element of grad_tensor belongs to whichever tensor supplied the element in that position. This function returns (up to) two gradients derived from grad_tensor split by the predicate: where(predicate, grad_tensor, 0) and where(predicate, 0, grad_tensor). The values of input_tensor_b and input_tensor_c do not affect the result.

Parameters:

  • grad_tensor (ttnn.Tensor)the gradient with respect to the output of the forward ttnn.where call.

  • input_tensor_a (ttnn.Tensor)the predicate tensor of the forward ttnn.where call, containing only 0’s or 1’s.

  • input_tensor_b (ttnn.Tensor)the true_value tensor of the forward ttnn.where call.

  • input_tensor_c (ttnn.Tensor)the false_value tensor of the forward ttnn.where call.

Keyword Arguments:

  • memory_config (ttnn.MemoryConfig, optional)memory configuration for the operation. Defaults to None.

  • are_required_outputs (List[bool], optional)which output gradients to compute. A False entry skips that gradient, and the returned list holds None in its place. Defaults to [True, True].

  • input_a_grad (ttnn.Tensor, optional)preallocated tensor to hold the first of the two results (corresponding to input_tensor_b). Defaults to None.

  • input_b_grad (ttnn.Tensor, optional)preallocated tensor to hold the second of the two results (corresponding to input_tensor_c). Defaults to None.

Returns:

List of ttnn.Tensorthe gradients with respect to input_tensor_b and input_tensor_c, or None if the corresponding entry in are_required_outputs is False.

Note

Supported dtypes and layouts:

Dtypes

Layouts

BFLOAT16, BFLOAT8_B

TILE

Example

# Create the forward inputs, with a predicate of 0's and 1's, and a gradient tensor
grad_tensor = ttnn.from_torch(
    torch.tensor([[1, 2], [3, 4]], dtype=torch.bfloat16), layout=ttnn.TILE_LAYOUT, device=device
)
predicate = ttnn.from_torch(
    torch.tensor([[1, 0], [1, 0]], dtype=torch.bfloat16), layout=ttnn.TILE_LAYOUT, device=device
)
true_value = ttnn.from_torch(
    torch.tensor([[5, 6], [7, 8]], dtype=torch.bfloat16, requires_grad=True), layout=ttnn.TILE_LAYOUT, device=device
)
false_value = ttnn.from_torch(
    torch.tensor([[9, 10], [11, 12]], dtype=torch.bfloat16, requires_grad=True),
    layout=ttnn.TILE_LAYOUT,
    device=device,
)

# Perform the where backward operation. Each gradient depends only on grad_tensor and the predicate,
# giving [[1, 0], [3, 0]] for true_value and [[0, 2], [0, 4]] for false_value
output = ttnn.where_bw(grad_tensor, predicate, true_value, false_value)
logger.info(f"Where Backward result: {output}")