Searched refs:TpuMeshStateInterface (Results 1 – 7 of 7) sorted by relevance
33 class TpuMeshStateInterface : public tensorflow::ResourceBase {35 explicit TpuMeshStateInterface(XLA_TpuMeshState* handle) in TpuMeshStateInterface() function39 ~TpuMeshStateInterface() override { in ~TpuMeshStateInterface()45 static TpuMeshStateInterface* Create() { in Create()46 return new TpuMeshStateInterface(OpsApiFn()->TpuMeshState_CreateFn()); in Create()
43 tpu::TpuMeshStateInterface** state) { in GetTpuMeshStateInterface()130 OP_REQUIRES_OK(ctx, DeleteIfExists<tpu::TpuMeshStateInterface>( in Compute()133 auto* tpu_mesh = tpu::TpuMeshStateInterface::Create(); in Compute()190 tpu::TpuMeshStateInterface* mesh_state; in Compute()240 OP_REQUIRES_OK(ctx, DeleteIfExists<tpu::TpuMeshStateInterface>( in Compute()262 OP_REQUIRES_OK(ctx, DeleteIfExists<tpu::TpuMeshStateInterface>( in Compute()264 auto* mesh_state_interface = tpu::TpuMeshStateInterface::Create(); in Compute()
40 const TpuMeshStateInterface& mesh_state);
121 const TpuMeshStateInterface* mesh_state,132 const TpuMeshStateInterface* mesh_state, in LookupPersistentCompilationCacheAndFillCaches()
109 const TpuMeshStateInterface& mesh_state) { in CreateCompilationCacheKey()
593 const TpuMeshStateInterface* mesh_state, in CompileLocallyAndFillHostCache()635 TpuMeshStateInterface* mesh_state; in ComputeInternal()
92 tensorflow::tpu::TpuMeshStateInterface* mesh_state, int num_replicas, in CompilationCacheKey()158 tensorflow::tpu::TpuMeshStateInterface* mesh_state; in Compute()