Torch cond operator, python dispatch, pyoperator

[voznesenskym]

Fixes #ISSUE_NUMBER

[facebook-github-bot]

🔗 Helpful links

• 🧪  See artifacts and rendered test results at hud.pytorch.org/pr/83154
• ✖️ Python docs build was skipped
• ✖️ C++ docs build was skipped
• ❓Need help or want to give feedback on the CI? Visit our office hours

✅ No Failures (3 Pending)

As of commit 670f3d2 (more details on the Dr. CI page):

Expand to see more

---

💚 💚 Looks good so far! There are no failures yet. 💚 💚

---

This comment was automatically generated by Dr. CI (expand for details).

Please report bugs/suggestions to the (internal) Dr. CI Users group.

Click here to manually regenerate this comment.

[voznesenskym]

Massive credit to @zou3519's work, this just keeps it moving

[voznesenskym]

Currently not working. pred is captured.

Which means that:

result_true = graph.forward(x, torch.tensor(True))
print("True:", result_true)
result_false = graph.forward(x, torch.tensor(False))
print("False:", result_false)

Have identical results. Need to figure this out. true_fn and false_fn parts are traced fine and work, changing the bool tensor passed to make_fx changes the result, when we really should be capturing it symbolically.

Need to figure out why those inputs arent going to conditional, but some tensor consts are getting created in their stead.

[zou3519]

I'm not sure this is going to fly, if we're doing real tracing (make_fx(f, mode='real')), it is supposed to end up executing real operations on the Tensor.

If we're tracing a function that calls cond and then passes the output of cond to another pytorch operation (like torch.sum), then torch.sum is going to receive this dummy Tensor and then will actually execute out.sum(). This will be a problem if the other operation is something that does rely on the output to be a real Tensor.

Here's an hacky idea that might work:

• create new_true_fn. new_true_fn executes true_fn and then stores the output of true_fn somewhere.

saved_true_out = []
def new_true_fn(*operands):
  out = true_fn(operands)
  saved_out.append(out)
  return out

• Do true_graph = get_isolated_graphmodule(new_true_fn, operands, {}) (and ditto for new_false_fn).
• Now, instead of out = torch.zeros([]), we can return either out = saved_true_out[0] or out = saved_false_out[0].
• Furthermore, we can compare saved_true_out and saved_false_out for if "all the properties match" (though maybe that is what your checks above are doing, I'm not familiar with the code).

[voznesenskym]

We might as well just run the real thing without all the extra pizzaz around new_true_fn though. Just uncomment my commented code.

Will change.

Furthermore, we can compare saved_true_out and saved_false_out for if "all the properties match" (though maybe that is what your checks above are doing, I'm not familiar with the code).

We already do that above, yes.

[zou3519]

Works for now, but I'm not sure it's comprehensive. E.g. TensorMetadata doesn't seem to have layout (torch.strided, torch.sparse_coo).

Also nit: don't use parenthesis while calling assert: assert true_out.meta == false_out.meta. The reason is that it is easy to accidentally turn the parenthesis into a tuple (e.g. assert (true_out.meta == false_out.meta,) and assert on a tuple checks that it has at least one element

[voznesenskym]

Meta is {'tensor_meta': TensorMetadata(shape=torch.Size([4]), dtype=torch.float32, requires_grad=False, stride=(1,), memory_format=torch.contiguous_format, is_quantized=False, qparams={})} for now - we can always add more to meta, right?

[malfet]

One needs to add empty dispatch/__init__.py otherwise it would not be packaged into neither wheel nor conda

[voznesenskym]

@pytorchbot merge

[pytorchmergebot]

@pytorchbot successfully started a merge job. Check the current status here.
The merge job was triggered without a flag. This means that your change will be merged once all checks on your PR have passed (ETA: 0-4 Hours). If this is not the intended behavior, feel free to use some of the other merge options in the wiki.
Please reach out to the PyTorch DevX Team with feedback or questions!

[pytorchmergebot]

Merge failed
Reason: View failures on hud. Refusing to merge as mandatory check(s) Lint, pull failed for rule Core Maintainers.
Raised by workflow job

[voznesenskym]

@pytorchbot merge

[pytorchmergebot]

@pytorchbot successfully started a merge job. Check the current status here.
The merge job was triggered without a flag. This means that your change will be merged once all checks on your PR have passed (ETA: 0-4 Hours). If this is not the intended behavior, feel free to use some of the other merge options in the wiki.
Please reach out to the PyTorch DevX Team with feedback or questions!

[github-actions]

Hey @voznesenskym.
You've committed this PR, but it does not have both a 'release notes: ...' and 'topics: ...' label. Please add one of each to the PR. The 'release notes: ...' label should represent the part of PyTorch that this PR changes (fx, autograd, distributed, etc) and the 'topics: ...' label should represent the kind of PR it is (not user facing, new feature, bug fix, perf improvement, etc). The list of valid labels can be found here for the 'release notes: ...' and here for the 'topics: ...'.
For changes that are 'topic: not user facing' there is no need for a release notes label.