Enable ConcreteFunction XLA JIT

Open
#347 10 comments 0 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

Assessment

Difficulty
4/5
Estimated time
3-5 days
Newbie friendliness
25/100
Issue type
Feature
Clarity
Needs clarification
Activity status
Stale
Tech stack
java
Domain
api, machine-learning

Research direction

Start with tensorflow-core/tensorflow-core-api/src/main/java/org/tensorflow/ConcreteFunction.java around line 475, then compare the behavior with tensorflow/python/eager/def_function.py around line 710. Review TensorFlow issue 50458 to determine whether the reported error is resolved. Done means ConcreteFunction XLA JIT can be enabled without that error and the behavior is verified.

Written by the indexing model from the issue text.

Description

This is mostly a reminder for myself, but also documentation if someone else ends up doing it.

After reading through the python code (see here), I found that to force function JITing you need to set the _XlaMustCompile and _noinline attributes to true, as is done here.

However, this results in an error. I didn't run it down at the time, but it was reported in https://github.com/tensorflow/tensorflow/issues/50458. Once that is fixed we should try this again.

Dominant language
Java
Stars
928
Forks
227
No merged PRs in 30d

Contributor guide

Open the contributing guide

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. Open a pull request that references the issue number.

More from tensorflow/java

All issues in tensorflow/java

Similar issues

More Java issues

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.