google / google/flax

How to use max_pool with xla_gpu_deterministic_ops?

Open
#4,233 0 comments 0 reactions 0 assignees View on GitHub
Dominant language
Jupyter Notebook
Stars
7.3k
Forks
833
Avg merge
5h 11m
Merged PRs (30d)
5

Description

I got the following error when I used `max_pool` in my module with `xla_gpu_deterministic_ops` turned on:

```
jaxlib.xla_extension.XlaRuntimeError: UNIMPLEMENTED: HLO instruction select-and-scatter.401.0 does not have a deterministic implementation, but run-to-run determinism is required by --xla_gpu_deterministic_ops.
```

Does it mean that I cannot use `max_pool` with `xla_gpu_deterministic_ops` turned on (so my code will not be reproducible)? Is there any workaround?

Contributor guide

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.