1

I have a pair of tensors S and T of dimensions (s1,...,sm) and (t1,...,tn) with si < ti. I want to specify a list of indices in each dimensions of T to "embed" S in T. If I1 is a list of s1 indices in (0,1,...,t1) and likewise for I2 up to In, I would like to do something like T.select(I1,...,In)=S that will have the effect that now T has entries equal to the entries of S over the indices (I1,...,In). for example

`S=
[[1,1],
[1,1]]

T=
[[0,0,0],
[0,0,0],
[0,0,0]]

T.select([0,2],[0,2])=S

T=
[[1,0,1],
[0,0,0],
[1,0,1]]`
kmario23
  • 57,311
  • 13
  • 161
  • 150
Vincent L.
  • 123
  • 5

1 Answers1

4

If you're flexible with using NumPy only for the indices part, then here's one approach by constructing an open mesh using numpy.ix_() and using this mesh to fill-in the values from the tensor S. If this is not acceptable, then you can use torch.meshgrid()

Below is an illustration of both approaches with descriptions interspersed in comments.

# input tensors to work with
In [174]: T 
Out[174]: 
tensor([[0, 0, 0],
        [0, 0, 0],
        [0, 0, 0]])

# I'm using unique tensor just for clarity; But any tensor should work.
In [175]: S 
Out[175]: 
tensor([[10, 11],
        [12, 13]])

# indices where we want the values from `S` to be filled in, along both dimensions
In [176]: idxs = [[0,2], [0,2]]

Now we will leverage np.ix_() or torch.meshgrid() to generate an open mesh by passing in the indices:

# mesh using `np.ix_`
In [177]: mesh = np.ix_(*idxs)

# as an alternative, we can use `torch.meshgrid()` 
In [191]: mesh = torch.meshgrid([torch.tensor(lst) for lst in idxs])

# replace the values from tensor `S` using basic indexing
In [178]: T[mesh] = S 

# sanity check!
In [179]: T 
Out[179]:
tensor([[10,  0, 11],
        [ 0,  0,  0],
        [12,  0, 13]])
kmario23
  • 57,311
  • 13
  • 161
  • 150