🔒 Passing optional keyword arguments to a function using `tree_map`

[inland nymph]

Hello,

This is mainly a Python question even thought I use PyTorch. In PyTorch we can define tensors and then do mathematical operations on them. We can define our own tensors and overwrite its behaviour. I'm overwriting a bunch of mathematical operations like e.g. how two of my custom tensors are being added. Writing your own tensor in PyTorch is done as a wrapper class where the data is stil hold by a normal pytorch tensor. So any function you don't define can still be called with the usual implementation.

To achieve this, we have three functions: (this implementation is faulty I think)

from torch.utils._pytree import tree_map

def unwrap(t):
    if isinstance(t, cls):
        return t.elem, t.max_grad, t.min_grad
    else:
        return t

def wrap(t, max_grad=None, min_grad=None):
    if isinstance(t, torch.Tensor) and not isinstance(t, cls):
        return cls(t, max_grad=max_grad, min_grad=min_grad, verbose=cls.verbose)
    else:
        return t

def run_with_usual_semantic():
    args_org = tree_map(unwrap, args)
    args_ = (args_org[0][0], *args_org[1:])
    kwargs_ = tree_map(unwrap, kwargs)
    res = func(*args_, **kwargs_)
    
    try:
        res = tree_map(wrap, res, args_org[0][1], args_org[0][2])
    except Exception as e:
        print("Error:", e)
        breakpoint()
    return res

• cls represents my own custom tensor subclass.
• t.elem represents the torch.tensor element that holds the actual data
• t.max_grad and t.min_grad are two additional optional tensors that I sometimes define when overwriting functions. The normal implementations of an functionof PyTorch does not know about these properties. If we call a function with the usual semantics i.e. the default implementation, we always call it with t.elem

[proper kestrel]

@inland nymph

Python help channel opened

Remember to:

• Ask your Python question, not if you can ask or if there's an expert who can help.
• Show a code sample as text (rather than a screenshot) and the error message, if you've got one.
• Explain what you expect to happen and what actually happens.

:warning: Do not pip install anything that isn't related to your question, especially if asked to over DMs.

Closes after a period of inactivity, or when you send !close.

[inland nymph]

So to summarize: If we call a function on our custom tensor, the function will always be applied to torch.tensor which is why we wrap and unwrap the data. E.g.

x = MyCustomTensor(torch.tensor([2]))
y = MyCustomTensor(torch.tensor([3]))
z = x + y

When we add x and y, we would unwrap both and actually add torch.tensor([2]) and torch.tensor([3]), the result would be torch.tensor([5]), we would wrap it and return MyCustomTensor(torch.tensor([5])).

Now there are a ton of different functions one can call with different signatures, so we use tree_map() to apply wrap and unwrap to each element.

Now as I said, sometimes MyCustomTensor might contain two extra elements: max_grad and min_grad and I'm not sure how I should handle those. Since they are optional, the default value of them is None.

So I had to change the three functions above. I think unwrap is fine, since the two properties always exist, we just return a tuple (t.elem, t.max_grad, t.min_grad) but with that, the problems begin.

Assume we call some function on our custom tensor that isn't implemented yet and we end up calling run_with_usual_semantic().

First we doargs_org = tree_map(unwrap, args) which gives us args_org = ((tensor(1.), None, None), [2], 0, 1). (Note that [2], 0 and 1 are just some random arguments of whatever function we call.)

We also get kwargs_ = tree_map(unwrap, kwargs) but we can ignore that for now. Now the function that is being called is provided as func. Since func is a normal torch function, it doesn't know about max_grad and min_grad, meaning the actual tuple we have to pass is (tensor(1.) [2], 0, 1). That's what we do in the line args_ = (args_org[0][0], *args_org[1:])

Note: I assume here that the first element args_org is always of type MyCustomTensor i.e. a tuple of 3 but that might not always be the case. We can ignore that for now.

Next we get the result: res = func(*args_, **kwargs_) and now we have to wrap it again.

I'm confused how to do that. Because now we have to add max_grad and min_grad again. That's what I try to do here: res = tree_map(wrap, res, args_org[0][1], args_org[0][2]) but I think this is wrong.

Here res = tensor([0., 1.]). So we could call tree_map(wrap, res) and get MyCustomTensor(tensor([0., 1.])) BUT now we lost max_grad and min_grad and I don't know how to pass those optiona lkeyword arguments to wrap using tree_map.

Any help is highly

[inland nymph]

I used partial from functols to prepopulate the keyword arguments.

.close

[sleek plover]

Did you mean:

[inland nymph]

!close

[proper kestrel]

Python help channel closed

This help channel has been closed and it's no longer possible to send messages here. If your question wasn't answered, feel free to create a new post in #1035199133436354600. To maximize your chances of getting a response, check out this guide on asking good questions.