ttnn.assign_bw

ttnn.assign_bw() → None

Returns the gradient of assign operation.

Example

# Create gradient and input tensors for assignment backward
grad_tensor = ttnn.from_torch(
    torch.tensor([[1, 2], [3, 4]], dtype=torch.bfloat16), layout=ttnn.TILE_LAYOUT, device=device
)
tensor1 = ttnn.from_torch(
    torch.tensor([[1, 2], [3, 4]], dtype=torch.bfloat16, requires_grad=True), layout=ttnn.TILE_LAYOUT, device=device
)
output = ttnn.assign_bw(grad_tensor, tensor1)
logger.info(f"Assignment backward result: {output}")

# Create second input tensor for assignment backward
tensor2 = ttnn.from_torch(
    torch.tensor([[1, 2], [3, 4]], dtype=torch.bfloat16, requires_grad=True), layout=ttnn.TILE_LAYOUT, device=device
)
output = ttnn.assign_bw(grad_tensor, tensor1, tensor2)
logger.info(f"Assignment backward with two inputs result: {output}")