Update build in Dense and EinsumDense for QuantizedDTypePolicy

[james77777778]

This PR updates the logic in build to directly add int8 weights if dtype_policy is a QuantizedDTypePolicy

Additionally, I have moved the logic about quantization to the bottom of the class.

[codecov-commenter]

Codecov Report

Attention: Patch coverage is 80.58252% with 20 lines in your changes are missing coverage. Please review.

Project coverage is 75.85%. Comparing base (ff28c35) to head (11ef4c4).
Report is 7 commits behind head on master.

Files	Patch %	Lines
keras/layers/core/einsum_dense.py	79.31%	5 Missing and 7 partials ⚠️
keras/layers/core/dense.py	82.22%	1 Missing and 7 partials ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##           master   #19347      +/-   ##
==========================================
+ Coverage   75.83%   75.85%   +0.02%     
==========================================
  Files         367      367              
  Lines       40371    40408      +37     
  Branches     7853     7861       +8     
==========================================
+ Hits        30614    30652      +38     
+ Misses       8065     8061       -4     
- Partials     1692     1695       +3     

Flag	Coverage Δ
keras	75.71% <80.58%> (+0.02%)	⬆️
keras-jax	60.14% <80.58%> (+0.03%)	⬆️
keras-numpy	54.41% <63.10%> (+0.02%)	⬆️
keras-tensorflow	61.31% <80.58%> (+0.03%)	⬆️
keras-torch	60.40% <80.58%> (+0.03%)	⬆️

Flags with carried forward coverage won't be shown. Click here to find out more.

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

[fchollet]

Thanks for the PR!

[fchollet]

It should default to 1, not 0, otherwise the layer's output would always be 0.

Remind me, why does it need to be a variable? Could it just be a constant?
And does it need to have shape (1, kernel_shape[1])? Could it be a scalar?

[james77777778]

It should default to 1, not 0, otherwise the layer's output would always be 0.

I've changed the defaults from "zeros" to "ones". However, it should be considered a dummy initialization because the quantized weights must be loaded from a pretrained model.

why does it need to be a variable? Could it just be a constant?

It would be reasonable to allow kernel_scale to be a variable. We can easily save and load it, and even modify it with the current APIs.

And does it need to have shape (1, kernel_shape[1])? Could it be a scalar?

I've changed the shape of kernel_scale in Dense to (self.units,) representing a 1D vector

Now, the current implementation is the same as google/gemma_pytorch:
https://github.com/google/gemma_pytorch/blob/main/gemma/model.py#L112-L121

            self.weight_scaler = nn.Parameter(torch.Tensor(out_features))

[fchollet]

LGTM, thank you!