vmap support for torch.tril, torch.triu

[zou3519]

Tasks

Add vmap support for the following PyTorch operations. That is, each one needs a batching rule.

• torch.tril
• torch.triu

These already have batching rules implemented, but the batching rules are incorrect.

Expected behavior

Currently, when one does vmap over them using a 2D input, it succeeds:

import torch
from functorch import vmap
x = torch.randn(32, 3)
y = vmap(torch.triu)(x)

This should actually raise an exception. The mental model here is that vmap should be equivalent to performing a for-loop over the arguments. The following expression does raise an exception, so vmap should as well:

results = []
for xi in x:
  y = torch.triu(xi)
  results.append(y)

Read this!

See this note for more context https://github.com/pytorch/pytorch/blob/master/functorch/writing_batching_rules.md

If you're new to developing PyTorch and/or function transforms, we would recommend reading through https://github.com/pytorch/pytorch/wiki/Core-Frontend-Onboarding#unit-8-function-transforms-optional

For this issue, you'll likely need to implement something similar to the vmap support for torch.clone:

• pytorch/aten/src/ATen/functorch/BatchRulesUnaryOps.cpp

  Line 95 in 06bdd49

  VMAP_SUPPORT(clone, clone_batch_rule);

Testing

pytest test/functorch/test_vmap.py -v -k "test_op_has_batch_rule_tril" should pass (currently, it is an expected failure). Ditto for triu.

cc @Chillee @samdow @soumith

[benjamintli]

ill give this a shot

[sunmorgan]

ill take this in

[sunmorgan]

@zou3519 It seems like torch.tril and torch.triu uses the "VARIADIC_BDIMS_BATCH_RULE" definition from BatchRulesHelper.h shown below:

pytorch/aten/src/ATen/functorch/BatchRulesHelper.h

Line 109 in 0a00858

#define VARIADIC_BDIMS_BATCH_RULE(fn) SINGLE_ARG(\

pytorch/aten/src/ATen/functorch/BatchRulesViews.cpp

Line 542 in df4b3b1

VMAP_SUPPORT(tril, VARIADIC_BDIMS_BATCH_RULE(ATEN_FN(tril)));

pytorch/aten/src/ATen/functorch/BatchRulesViews.cpp

Line 543 in df4b3b1

VMAP_SUPPORT(triu, VARIADIC_BDIMS_BATCH_RULE(ATEN_FN(triu)));

Then should the definition be removed? It doesn't look like the method is being used anywhere else

Also, when running tests on test/functorch/test_vmap.py, it outputs the module not found error "torch.testing._internal.autograd_function_db import autograd_function_db"

OS: Mac
Environment: Anaconda
Downloaded from source