Skip to content
Sign in
Success

Changes

Summary

  1. Fix OpenMP GPU host_data conversion (details)
  2. AMD/ROCM Support for Jax (#1376) (details)
  3. Prune inherited JAX pytests and remove binding dependencies (#1379) (details)
Commit e83e7d62179dd6c710b77216aea55df37600f820 by Gunther Huebler
Fix OpenMP GPU host_data conversion
The file was modifiedutilities/convert_acc_to_omp.bash (diff)
Commit da614f85c180af844a77e80ddd4b61a94ce716b9 by noreply
AMD/ROCM Support for Jax (#1376)

This does 2 things:

adds support for running on ROCM systems, this is the AMD GPU equivalent of cuda
organize specific-device backend setup files, e.g GPU on nvidia needs cuda jax libraries, Apple needs metal, AMD needs rocm
The only complication was the jax cholesky not working on AMD GPUs. Seems like a bug that's out of our control, so I left en environment flag to switch between the jax linalg call and a custom "portable" one. The portable one still passes the jax_vs_fortran comparison tests.


Commits
---
* Support native AMD ROCm JAX on CachyOS

* Keep ROCm crash fixes and remove Morrison rounding workarounds

* Separate JAX launcher and inspection backend policies
The file was modifiedclubb_jax/README.md (diff)
The file was addedclubb_jax/backends/cpu.py
The file was modifiedcompile.py (diff)
The file was modifiedBUILD.md (diff)
The file was addedclubb_jax/pytests/auto_llm_generated_pytests/test_backends.py
The file was addedclubb_jax/backends/__init__.py
The file was addedclubb_jax/requirements-rocm.txt
The file was modifiedtests/run_jax_vs_fortran_cases.py (diff)
The file was addedclubb_jax/backends/rocm.py
The file was modifiedclubb_jax/pytests/auto_llm_generated_pytests/test_runtime_info.py (diff)
The file was modifiedclubb_jax/pytests/auto_llm_generated_pytests/test_jax_cli_options.py (diff)
The file was addedtests/pytests/auto_llm_generated_pytests/test_compile_compiler_selection.py
The file was addedclubb_jax/pytests/auto_llm_generated_pytests/test_portable_cholesky.py
The file was modifiedclubb_jax/run_jax.py (diff)
The file was modifiedclubb_jax/src/CLUBB_core/matrix_operations.py (diff)
The file was addedclubb_jax/pytests/auto_llm_generated_pytests/test_rocm_runtime.py
The file was modifiedclubb_jax/runtime_info.py (diff)
The file was addedclubb_jax/backends/metal.py
The file was addedclubb_jax/backends/common.py
The file was addedclubb_jax/backends/cuda.py
Commit 350a98e42b5a193c8471cc67c288bce95e7f7ac0 by noreply
Prune inherited JAX pytests and remove binding dependencies (#1379)

This just prunes lots of outdated and unneccesary pytests. These are all LLM generated and I just never went through them before. Probably still needs cleanup.
The file was modifiedclubb_jax/pytests/auto_llm_generated_pytests/test_max_cubic_root.py (diff)
The file was modifiedclubb_jax/pytests/auto_llm_generated_pytests/test_inverse_hydrostatic.py (diff)
The file was removedclubb_jax/pytests/auto_llm_generated_pytests/test_adg2_driver.py
The file was removedclubb_jax/pytests/auto_llm_generated_pytests/test_mono_flux_limiter.py
The file was modifiedclubb_jax/pytests/auto_llm_generated_pytests/test_bivar_LL_covar_partial.py (diff)
The file was modifiedclubb_jax/pytests/auto_llm_generated_pytests/test_soil_vegetation.py (diff)
The file was modifiedclubb_jax/pytests/auto_llm_generated_pytests/test_cholesky_corr_mtx_approx.py (diff)
The file was modifiedclubb_jax/pytests/auto_llm_generated_pytests/test_cholesky_factor.py (diff)
The file was modifiedclubb_jax/pytests/auto_llm_generated_pytests/test_calc_xpthvp_terms.py (diff)
The file was modifiedclubb_jax/pytests/auto_llm_generated_pytests/test_precip_fraction.py (diff)
The file was modifiedclubb_jax/pytests/auto_llm_generated_pytests/test_w_up_in_cloud.py (diff)
The file was modifiedclubb_jax/pytests/auto_llm_generated_pytests/test_constants.py (diff)
The file was removedclubb_jax/pytests/auto_llm_generated_pytests/test_no_dead_functions.py
The file was modifiedclubb_jax/pytests/auto_llm_generated_pytests/test_rad_extended_atmosphere.py (diff)
The file was modifiedclubb_jax/pytests/auto_llm_generated_pytests/test_silhs_port_structure.py (diff)
The file was modifiedclubb_jax/pytests/auto_llm_generated_pytests/test_radiation_port_structure.py (diff)
The file was modifiedclubb_jax/pytests/auto_llm_generated_pytests/test_calc_wpxpyp_pdf.py (diff)
The file was modifiedclubb_jax/pytests/auto_llm_generated_pytests/test_differentiability.py (diff)
The file was modifiedclubb_jax/pytests/auto_llm_generated_pytests/test_kk_covariance_batch.py (diff)
The file was modifiedclubb_jax/pytests/auto_llm_generated_pytests/test_forcing_port_structure.py (diff)
The file was modifiedclubb_jax/pytests/auto_llm_generated_pytests/test_hydrometp2_zt.py (diff)
The file was modifiedclubb_jax/pytests/auto_llm_generated_pytests/test_spurious_source.py (diff)
The file was modifiedclubb_jax/pytests/auto_llm_generated_pytests/test_responder_params.py (diff)
The file was removedclubb_jax/pytests/auto_llm_generated_pytests/test_pos_definite.py
The file was modifiedclubb_jax/pytests/auto_llm_generated_pytests/test_pdf_params_init.py (diff)
The file was addedclubb_jax/tests/run_morrison_interface_test.py
The file was modifiedclubb_jax/pytests/auto_llm_generated_pytests/test_microphysics_integration.py (diff)
The file was removedclubb_jax/pytests/auto_llm_generated_pytests/test_routine_placement.py
The file was modifiedclubb_jax/pytests/auto_llm_generated_pytests/test_covar_KK_mixt_frac_linearity.py (diff)
The file was modifiedclubb_jax/pytests/auto_llm_generated_pytests/test_solver.py (diff)
The file was modifiedclubb_jax/pytests/auto_llm_generated_pytests/test_stats_writer.py (diff)
The file was modifiedclubb_jax/pytests/auto_llm_generated_pytests/test_clip_covar.py (diff)
The file was modifiedclubb_jax/pytests/auto_llm_generated_pytests/test_fill_holes_sliding_window.py (diff)
The file was removedclubb_jax/pytests/auto_llm_generated_pytests/test_calculate_thvm.py
The file was removedclubb_jax/pytests/auto_llm_generated_pytests/test_compute_sigma_sqd_w.py
The file was modifiedclubb_jax/pytests/auto_llm_generated_pytests/test_fill_holes_global.py (diff)
The file was modifiedclubb_jax/pytests/auto_llm_generated_pytests/test_calc_corr_w_hm_n.py (diff)
The file was removedclubb_jax/pytests/auto_llm_generated_pytests/test_xp2_xpyp_terms.py
The file was modifiedclubb_jax/pytests/auto_llm_generated_pytests/test_saturation.py (diff)
The file was removedclubb_jax/pytests/auto_llm_generated_pytests/test_diagnose_upxp.py
The file was modifiedclubb_jax/pytests/auto_llm_generated_pytests/test_ice_dfsn.py (diff)
The file was removedclubb_jax/pytests/auto_llm_generated_pytests/test_xpyp_term_ta_godunov.py
The file was modifiedclubb_jax/pytests/auto_llm_generated_pytests/test_validation_checks.py (diff)
The file was removedclubb_jax/pytests/auto_llm_generated_pytests/test_xpyp_term_ta_pdf.py
The file was modifiedclubb_jax/pytests/auto_llm_generated_pytests/test_hydromet_pdf_parameter.py (diff)
The file was modifiedclubb_jax/pytests/auto_llm_generated_pytests/test_compute_rtp2_from_chi.py (diff)
The file was removedclubb_jax/pytests/auto_llm_generated_pytests/test_compute_gamma_skw.py
The file was modifiedclubb_jax/pytests/auto_llm_generated_pytests/test_fill_holes_mean.py (diff)
The file was modifiedclubb_jax/pytests/auto_llm_generated_pytests/test_mpace_b_lba_tndcy.py (diff)
The file was modifiedclubb_jax/pytests/auto_llm_generated_pytests/test_jit_compile.py (diff)
The file was modifiedclubb_jax/pytests/auto_llm_generated_pytests/test_kk_local_means.py (diff)
The file was modifiedclubb_jax/pytests/auto_llm_generated_pytests/test_calc_comp_mu_sigma_hm.py (diff)
The file was removedclubb_jax/pytests/auto_llm_generated_pytests/test_windm_edsclrm_lhs.py
The file was modifiedclubb_jax/pytests/auto_llm_generated_pytests/test_index_mapping.py (diff)
The file was modifiedclubb_jax/pytests/auto_llm_generated_pytests/test_Nc_Ncn_eqns.py (diff)
The file was modifiedclubb_jax/pytests/auto_llm_generated_pytests/test_calc_comp_corrs_binormal.py (diff)
The file was modifiedclubb_jax/pytests/auto_llm_generated_pytests/test_diffusion.py (diff)
The file was removedclubb_jax/pytests/auto_llm_generated_pytests/test_mean_adv.py
The file was modifiedclubb_jax/pytests/auto_llm_generated_pytests/test_grid_operators.py (diff)
The file was modifiedclubb_jax/pytests/auto_llm_generated_pytests/test_T_in_K.py (diff)
The file was modifiedclubb_jax/pytests/auto_llm_generated_pytests/test_bivar_LL_mean_const.py (diff)
The file was modifiedclubb_jax/pytests/auto_llm_generated_pytests/test_limits_f_x_responder.py (diff)
The file was modifiedclubb_jax/pytests/auto_llm_generated_pytests/test_pressure_sounding_z.py (diff)
The file was modifiedclubb_jax/pytests/auto_llm_generated_pytests/test_quadrivar_NNLL_covar_eq_swap.py (diff)
The file was modifiedclubb_jax/pytests/auto_llm_generated_pytests/test_sfc_flux.py (diff)
The file was modifiedclubb_jax/pytests/auto_llm_generated_pytests/test_config_flags_complete.py (diff)
The file was modifiedclubb_jax/pytests/auto_llm_generated_pytests/test_remapping_module.py (diff)
The file was modifiedclubb_jax/pytests/auto_llm_generated_pytests/test_spec_hum_to_mixing_ratio.py (diff)
The file was removedclubb_jax/pytests/auto_llm_generated_pytests/test_cloud_drop_sed.py
The file was modifiedclubb_jax/pytests/auto_llm_generated_pytests/test_microphysics_error_paths.py (diff)
The file was modifiedclubb_jax/pytests/auto_llm_generated_pytests/test_remapping_ppm.py (diff)
The file was modifiedclubb_jax/pytests/auto_llm_generated_pytests/test_liq_water_path.py (diff)
The file was removedclubb_jax/pytests/auto_llm_generated_pytests/test_clip_skewness.py
The file was removedclubb_jax/pytests/auto_llm_generated_pytests/test_luhar_3d_driver.py
The file was modifiedclubb_jax/pytests/auto_llm_generated_pytests/test_pdf_binormal_moments.py (diff)
The file was modifiedclubb_jax/pytests/auto_llm_generated_pytests/test_trivar_NLL_mean_eq.py (diff)
The file was modifiedclubb_jax/pytests/auto_llm_generated_pytests/test_diagnose_correlations.py (diff)
The file was modifiedclubb_jax/pytests/auto_llm_generated_pytests/test_calc_corr_norm_cholesky.py (diff)
The file was modifiedclubb_jax/pytests/auto_llm_generated_pytests/test_component_precip_frac_specify.py (diff)
The file was removedclubb_jax/pytests/auto_llm_generated_pytests/test_mixing_length.py
The file was modifiedclubb_jax/pytests/auto_llm_generated_pytests/test_calc_ice_cloud_frac_component.py (diff)
The file was modifiedclubb_jax/pytests/auto_llm_generated_pytests/test_kk_upscaled_variances.py (diff)
The file was removedclubb_jax/pytests/auto_llm_generated_pytests/test_no_dead_imports.py
The file was removedclubb_jax/pytests/auto_llm_generated_pytests/test_sfc_varnce.py
The file was modifiedclubb_jax/pytests/auto_llm_generated_pytests/test_mirror_lower_triangular.py (diff)
The file was modifiedclubb_jax/pytests/auto_llm_generated_pytests/test_numerics_autodiff.py (diff)
The file was modifiedtests/run_pytests.sh (diff)
The file was removedclubb_jax/pytests/auto_llm_generated_pytests/test_cloud_correlate.py
The file was removedclubb_jax/pytests/auto_llm_generated_pytests/test_wp2_wp3_terms.py
The file was modifiedclubb_jax/pytests/auto_llm_generated_pytests/test_NNL_covar_const_limits.py (diff)
The file was modifiedclubb_jax/pytests/auto_llm_generated_pytests/test_rcm_sat_adj.py (diff)
The file was modifiedclubb_jax/pytests/auto_llm_generated_pytests/test_luhar_params.py (diff)
The file was modifiedclubb_jax/pytests/auto_llm_generated_pytests/test_pdf_moment_integrals.py (diff)
The file was modifiedclubb_jax/pytests/auto_llm_generated_pytests/test_calculate_w_responder_params.py (diff)
The file was modifiedclubb_jax/pytests/auto_llm_generated_pytests/test_ly93_pdf.py (diff)
The file was modifiedclubb_jax/pytests/auto_llm_generated_pytests/test_new_pdf.py (diff)
The file was removedclubb_jax/pytests/auto_llm_generated_pytests/test_penta_faithful.py
The file was removedclubb_jax/pytests/auto_llm_generated_pytests/test_wp2_wp3_lhs_terms.py
The file was addedclubb_jax/tests/run_standalone_no_bindings_test.py
The file was removedclubb_jax/pytests/auto_llm_generated_pytests/test_skx_func.py
The file was removedclubb_jax/pytests/auto_llm_generated_pytests/test_clip_variance.py
The file was modifiedclubb_jax/pytests/auto_llm_generated_pytests/test_sunray_sw.py (diff)
The file was removedclubb_jax/pytests/auto_llm_generated_pytests/test_set_sfc_value_of_flux_profiles.py
The file was removedclubb_jax/pytests/auto_llm_generated_pytests/test_coefs_semiimpl.py
The file was modifiedclubb_jax/pytests/auto_llm_generated_pytests/test_hydrometeor_mixed_moments.py (diff)
The file was modifiedclubb_jax/pytests/auto_llm_generated_pytests/test_cloud_feedback_sfclyr.py (diff)
The file was modifiedclubb_jax/README.md (diff)
The file was removedclubb_jax/pytests/auto_llm_generated_pytests/test_windm_edsclrm_rhs.py
The file was modifiedclubb_jax/pytests/auto_llm_generated_pytests/test_corr_varnce.py (diff)
The file was removedclubb_jax/pytests/auto_llm_generated_pytests/test_calc_wp3_on_wp2.py
The file was modifiedclubb_jax/pytests/auto_llm_generated_pytests/test_pdf_params_pack_roundtrip.py (diff)
The file was removedclubb_jax/pytests/auto_llm_generated_pytests/test_xm_wpxp_terms.py
The file was modifiedclubb_jax/pytests/auto_llm_generated_pytests/test_kk_autoconversion.py (diff)
The file was removedclubb_jax/pytests/auto_llm_generated_pytests/test_xm_correction_wpxp_cl.py
The file was modifiedclubb_jax/pytests/auto_llm_generated_pytests/test_mixed_moment_pdf_integrals.py (diff)
The file was modifiedclubb_jax/pytests/auto_llm_generated_pytests/test_quadrivar_covar_const_limits.py (diff)
The file was modifiedclubb_jax/pytests/auto_llm_generated_pytests/test_cos_solar_zen.py (diff)
The file was modifiedclubb_jax/pytests/auto_llm_generated_pytests/test_gabls3_night_stability.py (diff)
The file was modifiedclubb_jax/pytests/auto_llm_generated_pytests/test_trivar_NNL_covar_eq.py (diff)
The file was modifiedclubb_jax/pytests/auto_llm_generated_pytests/test_calc_roots.py (diff)
The file was modifiedclubb_jax/pytests/auto_llm_generated_pytests/test_sponge_damp_xp23.py (diff)
The file was modifiedclubb_jax/pytests/auto_llm_generated_pytests/test_setter_var_params.py (diff)
The file was modifiedclubb_jax/pytests/auto_llm_generated_pytests/test_adg1_adg2_responder_params.py (diff)
The file was removedclubb_jax/pytests/auto_llm_generated_pytests/test_smooth_heaviside.py
The file was modifiedjenkins_tests/clubb_jax/Jenkinsfile (diff)
The file was modifiedclubb_jax/pytests/auto_llm_generated_pytests/test_comp_corr_norm.py (diff)
The file was modifiedclubb_jax/pytests/auto_llm_generated_pytests/test_new_tsdadg_pdf.py (diff)
The file was modifiedLLM_prompts/pytest_workflow.md (diff)
The file was removedclubb_jax/pytests/auto_llm_generated_pytests/test_compute_uv_tndcy.py
The file was modifiedclubb_jax/pytests/auto_llm_generated_pytests/test_bivar_LL_covar_partial_rr_Nr.py (diff)
The file was modifiedclubb_jax/pytests/auto_llm_generated_pytests/test_penta_solver.py (diff)
The file was modifiedclubb_jax/pytests/auto_llm_generated_pytests/test_splatting_autodiff.py (diff)
The file was modifiedclubb_jax/pytests/auto_llm_generated_pytests/test_NNL_covar_vs_MM.py (diff)
The file was modifiedclubb_jax/pytests/auto_llm_generated_pytests/test_silhs_surface_schemes.py (diff)
The file was modifiedclubb_jax/pytests/auto_llm_generated_pytests/test_simple_rad.py (diff)
The file was modifiedclubb_jax/pytests/auto_llm_generated_pytests/test_pdf_utilities.py (diff)
The file was modifiedclubb_jax/pytests/auto_llm_generated_pytests/test_bivar_NL_mean_eq.py (diff)
The file was removedclubb_jax/pytests/auto_llm_generated_pytests/test_diagnostic_cache.py
The file was removedclubb_jax/pytests/auto_llm_generated_pytests/test_mean_vert_vel_up_down.py
The file was modifiedclubb_jax/pytests/auto_llm_generated_pytests/test_fill_holes_autodiff.py (diff)
The file was modifiedclubb_jax/pytests/auto_llm_generated_pytests/test_simple_rad_lba.py (diff)
The file was modifiedclubb_jax/pytests/auto_llm_generated_pytests/test_denorm_transform_corr.py (diff)
The file was modifiedclubb_jax/pytests/auto_llm_generated_pytests/test_numerical_check_setters.py (diff)
The file was modifiedclubb_jax/pytests/auto_llm_generated_pytests/test_pdf_integrals_all_mm.py (diff)
The file was modifiedclubb_jax/pytests/auto_llm_generated_pytests/test_calendar.py (diff)
The file was modifiedclubb_jax/pytests/auto_llm_generated_pytests/test_KK_Nrm_tendencies.py (diff)
The file was modifiedclubb_jax/pytests/auto_llm_generated_pytests/test_new_pdf_main.py (diff)
The file was modifiedclubb_jax/pytests/auto_llm_generated_pytests/test_transform_pdf_chi_eta_component.py (diff)
The file was modifiedclubb_jax/pytests/auto_llm_generated_pytests/test_stats_output_time.py (diff)
The file was modifiedclubb_jax/pytests/auto_llm_generated_pytests/test_advance_xp3_simplified.py (diff)
The file was addedclubb_jax/tests/run_pdf_monte_carlo_test.py
The file was removedclubb_jax/pytests/auto_llm_generated_pytests/test_calc_xprcp_component.py
The file was removedclubb_jax/pytests/auto_llm_generated_pytests/test_new_hybrid_pdf_main.py
The file was modifiedclubb_jax/pytests/auto_llm_generated_pytests/test_advance_helper_extras.py (diff)
The file was modifiedclubb_jax/pytests/auto_llm_generated_pytests/test_update_xp2_mc.py (diff)
The file was modifiedclubb_jax/pytests/auto_llm_generated_pytests/test_binormal_moments.py (diff)
The file was removedclubb_jax/pytests/auto_llm_generated_pytests/test_compute_xp3.py
The file was modifiedclubb_jax/pytests/auto_llm_generated_pytests/test_lba_sfclyr.py (diff)
The file was modifiedclubb_jax/pytests/auto_llm_generated_pytests/test_time_select.py (diff)
The file was removedclubb_jax/pytests/auto_llm_generated_pytests/test_standalone_jax.py
The file was removedclubb_jax/pytests/auto_llm_generated_pytests/test_hydrostatic.py
The file was modifiedclubb_jax/pytests/auto_llm_generated_pytests/test_interpolation.py (diff)
The file was removedclubb_jax/pytests/auto_llm_generated_pytests/test_pvertinterp.py
The file was removedclubb_jax/pytests/auto_llm_generated_pytests/test_bugsrad.py
The file was modifiedclubb_jax/pytests/auto_llm_generated_pytests/test_component_corr_ip.py (diff)
The file was modifiedtests/README.md (diff)
The file was modifiedclubb_jax/pytests/auto_llm_generated_pytests/test_calc_liquid_cloud_frac_component.py (diff)
The file was modifiedclubb_jax/pytests/auto_llm_generated_pytests/test_microphysics_port_structure.py (diff)
The file was modifiedclubb_jax/pytests/auto_llm_generated_pytests/test_close_luhar_pdf.py (diff)
The file was modifiedclubb_jax/pytests/auto_llm_generated_pytests/test_calc_wp4_pdf.py (diff)
The file was modifiedclubb_jax/pytests/auto_llm_generated_pytests/test_param_names.py (diff)
The file was removedclubb_jax/pytests/auto_llm_generated_pytests/test_NNL_covar_const_vs_MM.py