Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Update references to JAX's GitHub repo #752

Open
wants to merge 1 commit into
base: main
Choose a base branch
from
Open

Update references to JAX's GitHub repo #752

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
6 changes: 3 additions & 3 deletions aqt/jax/v2/aqt_dot_general.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -899,11 +899,11 @@ def _maybe_dequant(
if jax.local_devices()[0].platform == 'cpu':
# needed bet lax.dot_general(int4, int4) is illegal on cpu.
# TODO(aqt): Remove this platform check once
# https://github.com/google/jax/issues/19682 is fixed.
# https://github.com/jax-ml/jax/issues/19682 is fixed.
# TODO(yichizh): It's better to assert False here with the following msg
# msg = (
# 'lax.dot_general(int4, int4) is illegal on cpu:'
# ' https://github.com/google/jax/issues/19682. The simple workaround'
# ' https://github.com/jax-ml/jax/issues/19682. The simple workaround'
# ' is to upcast to int8, but in that case please directly set the'
# ' numerics bits to int8. Please contact the AQT team if you believe'
# ' the workaround is needed.'
Expand All @@ -920,7 +920,7 @@ def _maybe_dequant(
precision=lax.Precision.DEFAULT,
)
# TODO(lew): Do we have a correct precision above?
# Relevant: https://github.com/google/jax/issues/14022
# Relevant: https://github.com/jax-ml/jax/issues/14022
out = aqt_tensor.QTensor(
qvalue=out,
scale=[],
Expand Down
2 changes: 1 addition & 1 deletion aqt/jax/v2/aqt_dot_general_test.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -487,7 +487,7 @@ def test_fake_quant(
@parameterized.parameters([
dict(
# TODO(aqt): Change dlhs_bits to 4bit once
# https://github.com/google/jax/issues/19682 is fixed.
# https://github.com/jax-ml/jax/issues/19682 is fixed.
dg=config.config_v3(
fwd_bits=3,
dlhs_bits=6,
Expand Down
6 changes: 3 additions & 3 deletions aqt/jax_legacy/jax/README.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -7,10 +7,10 @@ quantization for convolution and matmul.

### Jax Libraries

- **quantization.quantized_dot**: [LAX.dot](https://github.com/google/jax/blob/f65a327c764406db45e95048dfe09209d8ef6d37/jax/_src/lax/lax.py#L632) with optionally quantized weights and activations.
- **quantization.quantized_dynamic_dot_general**: [LAX.dot general](https://github.com/google/jax/blob/f65a327c764406db45e95048dfe09209d8ef6d37/jax/_src/lax/lax.py#L667) with optionally quantized dynamic inputs.
- **quantization.quantized_dot**: [LAX.dot](https://github.com/jax-ml/jax/blob/f65a327c764406db45e95048dfe09209d8ef6d37/jax/_src/lax/lax.py#L632) with optionally quantized weights and activations.
- **quantization.quantized_dynamic_dot_general**: [LAX.dot general](https://github.com/jax-ml/jax/blob/f65a327c764406db45e95048dfe09209d8ef6d37/jax/_src/lax/lax.py#L667) with optionally quantized dynamic inputs.
- **quantization.quantized_sum**: Sums a tensor while quantizing intermediate accumulations.
- **quantization.dot_general_aqt**: Adds quantization to [LAX.dot_general](https://github.com/google/jax/blob/f65a327c764406db45e95048dfe09209d8ef6d37/jax/_src/lax/lax.py#L667) with option to use integer dot.
- **quantization.dot_general_aqt**: Adds quantization to [LAX.dot_general](https://github.com/jax-ml/jax/blob/f65a327c764406db45e95048dfe09209d8ef6d37/jax/_src/lax/lax.py#L667) with option to use integer dot.


### Flax Libraries
Expand Down
6 changes: 3 additions & 3 deletions aqt/jax_legacy/jax/quantization.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -703,7 +703,7 @@ def quantized_dot(
"""LAX dot with optionally quantized weights and activations.

Wraps LAX's `Dot
<https://github.com/google/jax/blob/f65a327c764406db45e95048dfe09209d8ef6d37/jax/_src/lax/lax.py#L632`_
<https://github.com/jax-ml/jax/blob/f65a327c764406db45e95048dfe09209d8ef6d37/jax/_src/lax/lax.py#L632`_
operator.

Args:
Expand Down Expand Up @@ -732,7 +732,7 @@ def quantized_dot(
ValueError: Shapes of 'act' and 'w' not compatible with quant_type.
"""
# This code was initially expanded from
# https://github.com/google/jax/blob/f65a327c764406db45e95048dfe09209d8ef6d37/jax/_src/lax/lax.py#L632
# https://github.com/jax-ml/jax/blob/f65a327c764406db45e95048dfe09209d8ef6d37/jax/_src/lax/lax.py#L632
# We keep the original return-value semantics of lax.dot, which this wraps. In
# particular, the type of the return value of quantized_dot is the same as the
# type of the inputs. That means that if the inputs are bfloat16, then the
Expand Down Expand Up @@ -1315,7 +1315,7 @@ def quantized_dynamic_dot_general(
"""LAX dot general with optionally quantized dynamic inputs.

Wraps LAX's `DotGeneral
<https://github.com/google/jax/blob/f65a327c764406db45e95048dfe09209d8ef6d37/jax/_src/lax/lax.py#L667`_
<https://github.com/jax-ml/jax/blob/f65a327c764406db45e95048dfe09209d8ef6d37/jax/_src/lax/lax.py#L667`_
operator.

Args:
Expand Down
Loading