Skip to content

vmap-incompatible inplace being used by existing functions #247

@Padarn

Description

@Padarn

Hi all,

I'm working on this PR #225, and am having some trouble fixing one of the tests which complains of

FAILED test/test_ops.py::TestOperatorsCPU::test_vmapvjp_clamp_cpu_float32 - RuntimeError: vmap: aten::logical_or_(self, *extra_args) is not possible because there exists a Tensor `other` in extra_args that has more elements than `self`. This happened due to `other` being vmapped over but `...

So I understand this error now, and where is it coming from: The backwards function for clamp ends up using logical_or_ in the way described. But, I don't understand what the ideal fix would be.

Some thoughts:

  • Change the backwards function for clamp so that it doesn't use the in place operation of logical_or?
  • Completely rewrite the batching rule for the backwards clamp so that it doesn't fallback into this path?

The second seems like a bad precedent, and the first seems a bit overkill.

So, any advice on how this kind of situation should be handled? Or is already handled by other parts of the codebase? (please feel free to point me to code rather than explaining)

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Type

    No type

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions