indexAccumulate python api

[jjsjann123]

Things done in this PR is to support embedding backward, which requires torch.index_put_(..., accumulate=True).

Stacked PRs:

• Adding IndexPutAccumulateOp #4063
• indexAccumulate python api #4066 <-- This PR

What this PR does:

• Added python API Tensor fd.ops.index_accumulate(Tensor acc, Tensor index, Tensor value

[github-actions]

Review updated until commit 4896186

Description

• Added Python API for index_put_accumulate

• Included test cases for index_put_accumulate

• Updated serialization and deserialization for IndexPutAccumulateOpRecord

• Added reference implementation for index_put_accumulate

---

Changes walkthrough 📝

	Relevant files
Enhancement	python_bindings.cpp Add Python API for index_put_accumulate                                    csrc/python_frontend/python_bindings.cpp Added index_put_accumulate function to nvf_ops Included documentation for index_put_accumulate +42/-0    fusion_record.cpp Register parser for IndexPutAccumulateOpRecord                      csrc/serde/fusion_record.cpp Registered parser for IndexPutAccumulateOpRecord +7/-0      fusion_record.h Define IndexPutAccumulateOpRecord                                                csrc/python_frontend/fusion_record.h Defined IndexPutAccumulateOpRecord struct +24/-0    fusion_cache.fbs Add IndexPutAccumulateOp to RecordType                                      csrc/serde/fusion_cache.fbs Added IndexPutAccumulateOp to RecordType enum +1/-0
Tests	opinfo_input_generators.py Add test case generator for index_put_accumulate                  tests/python/opinfo_input_generators.py Added index_put_accumulate_generator for test cases +19/-0    opinfos.py Add opinfo for index_put_accumulate                                            tests/python/opinfos.py Added index_put_accumulate_opinfo to shape_ops Included reference implementation index_put_accumulate_ref +28/-0

---

PR Reviewer Guide 🔍

Here are some key observations to aid the review process:

🧪 PR contains tests

⚡ Recommended focus areas for review Documentation The docstring for index_put_accumulate could be more detailed, especially regarding the restrictions and limitations compared to the full torch.index_put(..., accumulate=true). Accumulates values into a tensor at specified indices. This function performs a restricted version of `torch.index_put`. It adds the values from `value_tv` to the elements of `acc_tv` at the indices specified by `index_tv`. acc_tv: The tensor to accumulate into (in-place modification). index_tv: The tensor containing the indices. value_tv: The tensor containing the values to accumulate. Returns: An alias to the modified `acc_tv` tensor. Note: This is a restricted version and may not support all features of the full `torch.index_put(..., accumulate=true)` function. )doc"); Error Handling The IndexPutAccumulateOpRecord operator does not include any error handling. Consider adding checks for tensor dimensions and types to prevent runtime errors. void operator()(FusionState& fd) final { auto acc = fd.getFusionState(args_.at(0).index)->as<TensorView>(); auto index = fd.getFusionState(args_.at(1).index)->as<TensorView>(); auto value = fd.getFusionState(args_.at(2).index)->as<TensorView>(); auto output = indexPutAccumulate(acc, index, value); fd.setFusionState(outputs_.at(0).index, output); } Test Coverage The index_put_accumulate_generator only tests a single case. Consider adding more test cases with different tensor shapes, data types, and edge cases to ensure robustness. def index_put_accumulate_generator( op: OpInfo, dtype: torch.dtype, requires_grad: bool = False, **kwargs ): make_arg = partial( make_tensor, device="cuda", dtype=dtype, requires_grad=requires_grad ) make_index = partial(make_tensor, device="cuda", requires_grad=False) # vocab_size, hidden_size, seq_size cases = ((1024, 12, 300),) for vocab, hidden, seq in cases: for index_dtype in [torch.int, torch.long]: acc = make_arg((vocab, hidden)) index = make_index((seq,), low=0, high=vocab, dtype=index_dtype) value = make_arg((seq, hidden)) yield SampleInput(acc, index, value)

[jjsjann123]

marking this as draft to avoid accidental merge.
But this PR is good for review as-is.

[rdspring1]

Do you need to define void handle(IndexAccumulateOp* iaop) in csrc/python_frontend/translation.cpp for the python clone and segmentation features?

Otherwise, the PR looks good to me.

[rdspring1]

I'm trying to improve python user experience by adding a docstring to new functions.

Docstring generated by Gemini.

    m.def("index_accumulate", &indexAccumulate,
          py::arg("acc_tv"), py::arg("index_tv"), py::arg("value_tv"),
          R"(
        Accumulates values into a tensor at specified indices.

        This function performs a restricted version of `torch.index_put(..., accumulate=true)`.
        It adds the values from `value_tv` to the elements of `acc_tv` at the indices
        specified by `index_tv`.

        acc_tv: The tensor to accumulate into (in-place modification).
        index_tv: The tensor containing the indices.
        value_tv: The tensor containing the values to accumulate.

        Returns:
            A pointer to the modified `acc_tv` tensor.

        Note:
            This is a restricted version and may not support all features of the
            full `torch.index_put(..., accumulate=true)` function.
    )");

[jjsjann123]

Hahaha, thanks for the draft~~~ will add it in.

[jjsjann123]

!test

[rdspring1]

👍🏼

[jjsjann123]

!test

[jjsjann123]

!test