ROCm/jax

Pull Requests Commits

have loops module generate same-argument jaxprs for single-operand cond

froystig committed 5 years ago

efc1104c

transpose rule for single-operand cond primitive

froystig committed 5 years ago

30bb5fdd

partial evaluation rule for single-operand cond primitive

froystig committed 5 years ago

37bb6f09

fix cond-with-constants tests to use jax.numpy where needed

froystig committed 5 years ago

139c2a96

JVP rule for single-operand cond primitive

froystig committed 5 years ago

34d75c2e

batching rule for single-operand cond primitive

froystig committed 5 years ago

226b5821

merge constvars when forming cond branch jaxprs

froystig committed 5 years ago

fc4ab77b

translation rule for single-operand cond primitive

froystig committed 5 years ago

027a5390

bind a single-operand cond primitive and update jaxpr typechecks

froystig committed 5 years ago

97738be4

update uses of cond in lax control flow tests

froystig committed 5 years ago

28e698ed

take a single operand in lax.cond and deprecate the old calling convention

froystig committed 5 years ago

92176418

Update jaxlib version to 0.1.47 in README

skye committed 5 years ago

16cf8451

disable arr.view() test on TPU (#3089)

jakevdp committed 5 years ago

Verified 6bd01602

promote integer inputs to float in jnp.median() and jnp.quantile() (#3082)

jakevdp committed 5 years ago

Verified 777636af

Implement .view() method of jax.numpy arrays (#3073)

jakevdp committed 5 years ago

Verified 59aab016

Remove workaround for Mac linear algebra bug that is fixed in the minimum jaxlib version. (#3080)

hawkinsp committed 5 years ago

Verified 22d14fd7

Disable trapz test on TPU. (#3078)

hawkinsp committed 5 years ago

Verified 91d1e0dd

Make jnp.array convert empty list to DeviceArray (#3049)

sharadmv committed 5 years ago

Verified e9d33946

Change instances of onp to np and np to jnp (#3044)

joaogui1 committed 5 years ago

Verified d2f84d63

Disable check_type for trapz test due to test failures. (#3071)

hawkinsp committed 5 years ago

Verified cd966f28

Avoid recompilation of rolled loops in threefry2x32. (#3069)

tomhennigan committed 5 years ago

Verified abdf504e

Refactor aval_to_result_handler to take unsharded aval. (#3067)

skye committed 5 years ago

Verified 11760ca9

lower in_axes=None to XLA replication annotation (#3025)

jekbradbury committed 5 years ago

Verified 4a84b913

Add special case for integer scalars to jax.numpy.power. (#3066)

hawkinsp committed 5 years ago

Verified 88d3422a

Add lax implementation of np.trapz (#3042)

jakevdp committed 5 years ago

Verified 96fbfeee

Update jax version to 0.1.67 (#3065)

skye committed 5 years ago

Verified ef4debca

Import tpu_driver after xla_client (#3064)

skye committed 5 years ago

Verified 0d97c3ba

Fix error message in optimizers piecewise_constant. (#3061)

hawkinsp committed 5 years ago

Verified 8008aa9e

Uses jnp.square instead of power. (#3036)

Yusuke Oda committed 5 years ago

Verified ccb8d459

Adjusted lax.numpy.indices test for older versions of numpy (#3053)

gnecula committed 5 years ago

Verified 28bc4b75

Newer Older