Checks if a function intended for a specific platform can be executed on the current one.
(target_platform: str)
| 104 | |
| 105 | |
| 106 | def is_supported_platform(target_platform: str) -> bool: |
| 107 | """Checks if a function intended for a specific platform can be executed on the current one.""" |
| 108 | devices = jax.devices() |
| 109 | supported = all(device.platform == target_platform for device in devices) |
| 110 | if not supported: |
| 111 | logging.info( |
| 112 | "Skipping test for %s on %s", |
| 113 | target_platform, |
| 114 | [device.platform for device in devices], |
| 115 | ) |
| 116 | return supported |
| 117 | |
| 118 | |
| 119 | def is_supported_mesh_shape( |