Copying into a slice [Feature Request]

[athul-bos-semi]

In order to perform tensor manipulation operations (such as Image-to-column), copying into specific slices of a tensor is important. The current tt_lib.tensor.copy can copy from slices, but will only work if the destination is a full tensor of the same slice.

The solution I'd like

import ttnn
import tt_lib

import torch

ttnn.set_printoptions(profile="full")

device_id = 3
device = tt_lib.device.CreateDevice(device_id)

activations = torch.randn((1, 1, 1, 10), dtype=torch.bfloat16)
act = tt_lib.tensor.Tensor(activations)
act = act.to(device)

# columns = torch.zeros((1, 1, 1, 16), dtype=torch.bfloat16)
columns = torch.zeros((1, 1, 1, 10), dtype=torch.bfloat16)
col = tt_lib.tensor.Tensor(columns)
col = col.to(device)

# print(dir(act))
print(act)
print(col)

print("\n------ Copying ------\n")

# print(act[0,0,0].get_legacy_shape())
# print(col[0,0,0,:act.volume()].get_legacy_shape())
# print(tt_lib.tensor.copy(act[0,0,0], col[0,0,0,:act.volume()]))
# tt_lib.tensor.copy(act[0,0,0], col[0,0,0,:act.volume()])

print(act[:,:,:,:].get_legacy_shape())

print()
print(col[:,:,:,:].get_legacy_shape())
tt_lib.tensor.copy(act[:,:,:,:], col[:,:,:,:])
print(col)

print()
print(col.get_legacy_shape())
tt_lib.tensor.copy(act[:,:,:,:], col)
print(col)

tt_lib.device.CloseDevice(device)

It would be great if tt_lib.tensor.copy(act[:,:,:,:], col[:,:,:,:]) also works the same as tt_lib.tensor.copy(act[:,:,:,:], col), and consequently, copy into a tensor of different shape like tt_lib.tensor.copy(act[:,:,:,:], col[:,:,:,:act.volume()]) (the destination slice will still be the same shape.

Output of current code

Environment information: OS: Devcloud running on Ubuntu 20.04.4 LTS (Focal) Version of software: v0.46.0-rc7

[davorchap]

FYI @mbahnasTT

[davorchap]

@jliangTT , TM team to triage. fyi @tarafdarTT @yan-zaretskiy

[jliangTT]

@davorchap , prirotiy pelase?

[athul-bos-semi]

@davorchap @jliangTT Please consider this one of the Higher priority Ops, as Copying into slices make Tensor Manipulation becomes a whole lot easier in the long run

[athul-bos-semi]

@davorchap Is there any update on the priority for this Op?