Automatically add out kwarg & member & inplace_member

[mavenlin]

As per our discussion in #18,

• For functions that support out=None kwarg, we use a preset template in the decorator.
• Automatically set the member functions when they exist to avoid manually add them at two places. When we implement torch.add, torch.Tensor.add and torch.Tensor.add_ are automatically added. While it works for most of the functions, for some of the functions this can be turned off, e.g. torch.permute & torch.Tensor.permute has different signatures.

[samuela]

Thanks for putting this together @mavenlin ! Let's also write some automated tests for all these out=tensor things and Tensor.foo_ methods

[mavenlin]

To enable automatic tests of out, foo & foo_, I did some refactoring on the test code.
Summarized as follows

• To automatically test member functions, the tests are required to directly take the original torch.foo instead of lambda x: torch.foo(x, a=1, b=2). Therefore, the kwargs {a=1, b=2} needs to be passed instead of binding via lambda functions.
• Use tree_map to handle the output of t2j_function, so that functions like torch.unbind can be tested as it is, without needing to wrap with a lambda.
• Testing with test_t2j_function(lambda ... will still go through forward & backward tests like the original code, but will not be tested for out and member functions.

It is not in the best way I would like it though, let me know whatever suggestions you have.

[mavenlin]

The CI error seems to be from nix version before the checkout step, I'm not very familiar with that, @samuela could you help take a look?

[samuela]

yes, sorry about the ci. it should be fixed after b958825

[mavenlin]

Great, tests are passing after the ci fix. Let me know if you have further comments on this one.

[samuela]

looks like there are some merge conflicts. i'm also gonna have to meditate a bit on how exactly i want to design this. what u have is great, but i'm concerned we may eventually blow up the complexity of implements

[mavenlin]

looks like there are some merge conflicts. i'm also gonna have to meditate a bit on how exactly i want to design this. what u have is great, but i'm concerned we may eventually blow up the complexity of implements

The awkward part is that whether out_kwarg is added interferes with Torchishify_output. Adding member to torch.Tensor also requires the overridden torch function's name.

[mavenlin]

2025-08-19

• Explicit control what to test for each function
• Test gradient for all cases by reducing output to scalar.
• Fix gradient of batch_norm

[samuela]

Thanks @mavenlin !