Allow passing custom dataset adapters.

[Hilly12]

• Allow passing a custom dataset adapter - this is to enable enqueuing ID tensors to TPUEmbedding before converting them to Jax arrays and putting them on tensorcores.
• Small fix in compute_output_shape in embedding layer.

[codecov-commenter]

Codecov Report

Attention: Patch coverage is 50.00000% with 2 lines in your changes missing coverage. Please review.

Project coverage is 78.87%. Comparing base (5aa5f88) to head (f0eb570).
Report is 12 commits behind head on master.

Files with missing lines	Patch %	Lines
keras/src/trainers/data_adapters/__init__.py	33.33%	1 Missing and 1 partial ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##           master   #20349      +/-   ##
==========================================
+ Coverage   78.81%   78.87%   +0.06%     
==========================================
  Files         512      512              
  Lines       49063    49264     +201     
  Branches     9035     7953    -1082     
==========================================
+ Hits        38668    38859     +191     
- Misses       8530     8546      +16     
+ Partials     1865     1859       -6     

Flag	Coverage Δ
keras	78.74% <50.00%> (+0.07%)	⬆️
keras-jax	62.40% <50.00%> (+0.13%)	⬆️
keras-numpy	57.42% <50.00%> (+0.01%)	⬆️
keras-tensorflow	63.64% <50.00%> (+0.08%)	⬆️
keras-torch	62.37% <50.00%> (+0.11%)	⬆️

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.