(inp, pad_value, padding)
| 318 | |
| 319 | |
| 320 | def pad(inp, pad_value, padding): |
| 321 | # interior is used as dilated padding if it is not zero |
| 322 | assert isinstance( |
| 323 | pad_value, (int, float, bool, np.ndarray) |
| 324 | ), f"pad_value error {type(pad_value)}" |
| 325 | pad_value = HLOTensor(pad_value, dtype=inp.dtype) |
| 326 | |
| 327 | low, high, interior = [], [], [] |
| 328 | for p in padding: |
| 329 | assert len(p) == 3 |
| 330 | low.append(p[0]) |
| 331 | high.append(p[1]) |
| 332 | interior.append(p[2]) |
| 333 | |
| 334 | return HLOTensor( |
| 335 | hlo.PadOp( |
| 336 | inp.tensor, |
| 337 | pad_value.tensor, |
| 338 | ir_utils.dense_int_elements(low), |
| 339 | ir_utils.dense_int_elements(high), |
| 340 | ir_utils.dense_int_elements(interior), |
| 341 | ).result |
| 342 | ) |
| 343 | |
| 344 | |
| 345 | def xla_gather( |
no test coverage detected