Skip to content

[Relax][Frontend][TFLite] Add RANDOM_UNIFORM, RANDOM_STANDARD_NORMAL, and MULTINOMIAL#19473

Merged
tlopex merged 2 commits into
apache:mainfrom
Aharrypotter:tflite-multinomial-random-19412
May 1, 2026
Merged

[Relax][Frontend][TFLite] Add RANDOM_UNIFORM, RANDOM_STANDARD_NORMAL, and MULTINOMIAL#19473
tlopex merged 2 commits into
apache:mainfrom
Aharrypotter:tflite-multinomial-random-19412

Conversation

@Aharrypotter

@Aharrypotter Aharrypotter commented Apr 29, 2026

Copy link
Copy Markdown
Contributor

Summary

This PR adds support for the TFLite RANDOM_UNIFORM, RANDOM_STANDARD_NORMAL, and MULTINOMIAL operators in the Relax TFLite frontend, covering the items I claimed in #19412.

RANDOM_UNIFORM and RANDOM_STANDARD_NORMAL are lowered to seeded tvm.contrib.random calls with dynamic shape support. MULTINOMIAL is lowered by composing existing Relax ops around relax.op.multinomial_from_uniform.

Changes

Frontend

  1. Add converter registrations for RANDOM_UNIFORM, RANDOM_STANDARD_NORMAL, and MULTINOMIAL.
  2. Add shared helpers to:
    • parse TFLite RandomOptions seeds
    • convert shape tensors to Relax shape expressions for dynamic-shape random ops
  3. Lower RANDOM_UNIFORM to tvm.contrib.random.uniform.
  4. Lower RANDOM_STANDARD_NORMAL to tvm.contrib.random.normal.
  5. Lower MULTINOMIAL by:
    • applying R.nn.softmax to logits
    • generating seeded uniform samples
    • calling R.multinomial_from_uniform
    • reshaping the result to [batch_size, num_samples]

Runtime

  1. Extend src/runtime/contrib/random/random.cc to accept seeded calls for tvm.contrib.random.uniform and tvm.contrib.random.normal.
  2. Preserve compatibility with the existing unseeded calling convention.

Testing

All tests pass:

pytest tests/python/relax/test_frontend_tflite.py::test_random_uniform_dynamic_shape \
       tests/python/relax/test_frontend_tflite.py::test_random_standard_normal_dynamic_shape \
       tests/python/relax/test_frontend_tflite.py::test_multinomial_dynamic_num_samples -v

References

@gemini-code-assist gemini-code-assist Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Code Review

This pull request adds support for the MULTINOMIAL, RANDOM_STANDARD_NORMAL, and RANDOM_UNIFORM operators to the TFLite frontend in Relax. It also updates the random contrib runtime to support seeded operations and flexible argument handling. Review feedback suggests improving the seed combination logic to handle negative values, optimizing the thread-local random engine to prevent redundant re-seeding, and enhancing the maintainability of error messages and assertions.

Comment thread src/runtime/contrib/random/random.cc Outdated
Comment thread src/runtime/contrib/random/random.cc
Comment thread python/tvm/relax/frontend/tflite/tflite_frontend.py Outdated
Comment thread python/tvm/relax/frontend/tflite/tflite_frontend.py Outdated

@tlopex tlopex left a comment

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Some points you'd better fix:

  1. Fix the nightly E2E tests. Since TF and TVM use different RNGs, do not compare exact numerical values. Instead, compare statistical properties, verify shape/dtype/call counts, check TVM self-consistency, or just skip E2E for these ops entirely.

  2. Clean up the dead code in random.cc. Remove the unreachable output-first (try_cast) branches. If you need to keep them to defend against future callers, add a clear comment explaining this.

  3. Add comments in random.cc and the converter's docstring to explicitly state that this uses "stateless RNG" semantics (i.e., identical seeds will re-seed and produce identical results on every call).

  4. Address the data truncation in CombineSeeds. Either adjust the mixing logic to avoid collisions when the high 32 bits of negative numbers are truncated, or add a comment noting that this truncation is intentional.

  5. Add a one-line comment for the seed=0 and seed2=0 fallback to the global engine, clarifying that this aligns with the TF spec for non-deterministic seeds.

  6. Fix the dispatch dictionary ordering. Move MULTINOMIAL after MUL so it stays in alphabetical order.

  7. Check if other frontends (like ONNX or PyTorch) already have a helper function for the tensor to ShapeExpr conversion. If they do, reuse it to avoid duplicate boilerplate code.

  8. Print the IR for test_multinomial_dynamic_num_samples to confirm TFLite didn't constant-fold the scalar input during conversion. This ensures the dynamic branch is actually being tested.

  9. Add validation for the multinomial output_dtype so it only accepts int32 or int64, or add a comment explicitly stating that the converter forwards whatever integer type TFLite produces.

@Aharrypotter Aharrypotter force-pushed the tflite-multinomial-random-19412 branch from 31fd65e to 2da3404 Compare May 1, 2026 04:44
@Aharrypotter

Copy link
Copy Markdown
Contributor Author

Some points you'd better fix:

  1. Fix the nightly E2E tests. Since TF and TVM use different RNGs, do not compare exact numerical values. Instead, compare statistical properties, verify shape/dtype/call counts, check TVM self-consistency, or just skip E2E for these ops entirely.
  2. Clean up the dead code in random.cc. Remove the unreachable output-first (try_cast) branches. If you need to keep them to defend against future callers, add a clear comment explaining this.
  3. Add comments in random.cc and the converter's docstring to explicitly state that this uses "stateless RNG" semantics (i.e., identical seeds will re-seed and produce identical results on every call).
  4. Address the data truncation in CombineSeeds. Either adjust the mixing logic to avoid collisions when the high 32 bits of negative numbers are truncated, or add a comment noting that this truncation is intentional.
  5. Add a one-line comment for the seed=0 and seed2=0 fallback to the global engine, clarifying that this aligns with the TF spec for non-deterministic seeds.
  6. Fix the dispatch dictionary ordering. Move MULTINOMIAL after MUL so it stays in alphabetical order.
  7. Check if other frontends (like ONNX or PyTorch) already have a helper function for the tensor to ShapeExpr conversion. If they do, reuse it to avoid duplicate boilerplate code.
  8. Print the IR for test_multinomial_dynamic_num_samples to confirm TFLite didn't constant-fold the scalar input during conversion. This ensures the dynamic branch is actually being tested.
  9. Add validation for the multinomial output_dtype so it only accepts int32 or int64, or add a comment explicitly stating that the converter forwards whatever integer type TFLite produces.

Thanks for the detailed review. I addressed the two larger frontend/test points as follows:

For the tensor-to-ShapeExpr conversion, I kept the helper local to the TFLite frontend instead of introducing a dependency on another frontend's private helper. PR #19433 had already added the same tensor_to_shape + ShapeStructInfo boilerplate for dynamic FILL
dims, and this PR introduced the same pattern for RANDOM_UNIFORM / RANDOM_STANDARD_NORMAL shape tensors. I refactored that into _get_shape_expr_from_tensor(...) and now reuse it from:

  • convert_fill
  • convert_random_uniform
  • convert_random_standard_normal

I left the MULTINOMIAL dynamic num_samples path separate because it starts from a scalar tensor and explicitly reshapes it to [1] before binding multinomial_num_samples, so it is not exactly the same shape-tensor case.

For the nightly E2E tests, I removed exact TF-vs-TVM random value comparison. Since TF and TVM use different RNG implementations, the random tests now validate:

  • imported Relax IR structure
  • output shape matches TensorFlow
  • output dtype matches TensorFlow
  • TVM seeded self-consistency across repeated execution with the same inputs/seeds

test_multinomial_dynamic_num_samples also checks for both R.tensor_to_shape and multinomial_num_samples in the script output, so the dynamic branch is actually exercised instead of being constant-folded away.

@Aharrypotter

Copy link
Copy Markdown
Contributor Author

cc @tlopex

@tlopex tlopex left a comment

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

LGTM

@tlopex tlopex merged commit 71c634f into apache:main May 1, 2026
9 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants