I very much like this manuscript and the effort to create a new atmospheric LES with a modern, flexible programming paradigm JAX. The twofold benefit: GPU acceleration and automatic differentiation, are well explained and illustrated with the warm bubble test cases and performance tables. I believe this will be a valuable research tool for future investigations into the value of machine learning for parameterizations crossing the grey zone.

Some minor remarks:

• Line 57: There is a listing of the advantages of the hybrid machine learning approach to atmospheric sciences, but weather forecast accuracy is in my opinion not one of them. Purely data-driven DL models, such as Google's graphcast and ECMWF's AIFS nowadays exceed forecast accuracy of both physical and hybrid-ML models. So I would recommend to remove the 'accurate weather predictions' from this line.
• In line 195 the authors explain that the accumulated MSE forms the basis of the loss calculation for training the SGS. Are the authors concerned with a 'smearing' effect of the SGS tendency, could this be observed?
• In line 203 it is explained that water vapor is clipped to positive values after each correction time step. Could this be resolved with an extra loss term?
• I would drop the end of the first paragraph in section 5.2, "This can be further proven...". It comes across as over-explanatory and the last sentence is not proper English.
• line 336: the statement that running LEX on 1 GPU is as fast as running CM1 on 600 cores does not follow from the numbers presented, it would require a strong scaling study from CM1 to verify this, and if that scaling is not ideal, CM1 will not reach the LEX performance at all on 600 cores. So please indicate whether this statement follows from separate performance measurements of CM1.
• Reference: I believe a reference to "JAX-Fluids: A fully-differentiable high-order computational fluid dynamics solver for compressible two-phase flows" (https://doi.org/10.1016/j.cpc.2022.108527) should be added as it closely aligns with this work.

Some missing pieces in this paper:

• I believe the manuscript lacks a proper explanation of the applied boundary conditions. I suspect you are using periodic lateral boundaries and a sponge layer at the top? Pleaase elaborate in the theoretical section.
• I would love to see a short example of a more realistic emergent cumulus case, and especially whether the trained SGS parameterization from the warm bubble can be transferred to a cloudy atmosphere.
• The manuscript lacks an outlook with respect to missing components in LAX: microphysics, radiation, surface scheme. It should be mentioned in the article what the status of these elements are and how this limits the applicability of LAX.
• There is no multi-GPU benchmarks in the paper. I believe JAX with XLA can scale across many GPU's, does LAX also scale beyond a single device? Please elaborate in the performance section.
• There is no mentioning of hyperparameter choices or tuning thereof in the SGS training section. It could be nice to add a small exploration of this. 

Please find my comments in the attached pdf.