How can I map a tensor by a lookup table?

I’d like to make an autograd function which will map all values based on a lookup table where each input value bin is mapped to a fixed output value from a lookup table. The mapping table may also depend on a dimension index of the tensor.

What could be a way to achieve that?

(in numpy I would keep using something like searchsorted to find the input value index in the bins, and then return the output value at the index)