openxla / xla

A machine learning compiler for GPUs, CPUs, and ML accelerators
Apache License 2.0
2.39k stars 358 forks source link

PR #14158: Add kExp op to algsimp IsNonNegative #14249

Closed copybara-service[bot] closed 5 days ago

copybara-service[bot] commented 5 days ago

PR #14158: Add kExp op to algsimp IsNonNegative

Imported from GitHub PR https://github.com/openxla/xla/pull/14158

kExp result is always non-negative (similar to kAbs). We can add kExp to IsNonNegative() util function in algsimp

Test:

import jax.numpy as jnp                                                                                                                                                                                                                 
x = jnp.array([0.0, -0.0, jnp.inf, -jnp.inf, jnp.nan, 0.1, -0.1, 3, -3]) 

jnp.abs(x)
# Array([0. , 0. , inf, inf, nan, 0.1, 0.1, 3. , 3. ], dtype=float32)

jnp.exp(x)
# Array([ 1. , 1. , inf, 0. , nan, 1.105171, 0.9048374, 20.085537, 0.04978707], dtype=float32)

Adrian, what you think? @akuegel Copybara import of the project:

-- 401c3e90a1bac7ef4594ad5b6faee5fab4e21871 by Alexander Pivovarov pivovaa@amazon.com:

Add kExp op to algsimp IsNonNegative

and add IsNonNegative tests

Merging this change closes #14158

FUTURE_COPYBARA_INTEGRATE_REVIEW=https://github.com/openxla/xla/pull/14158 from apivovarov:exp_is_non_negative 401c3e90a1bac7ef4594ad5b6faee5fab4e21871